diff --git a/docs/api/attention.md b/docs/api/attention.md index 02464da8..f731fb70 100644 --- a/docs/api/attention.md +++ b/docs/api/attention.md @@ -7,7 +7,7 @@ what runs when you call `op(...)`. ## Multi-Head Attention -::: tileops.attention.MHADecodePagedWithKVCacheFwdOp +::: tileops.attention.MHAPagedFwdOp options: show_root_heading: true heading_level: 3 @@ -15,7 +15,7 @@ what runs when you call `op(...)`. ## Grouped-Query Attention -::: tileops.attention.GQABwdOp +::: tileops.attention.GQADenseBwdOp options: show_root_heading: true heading_level: 3 @@ -41,7 +41,7 @@ what runs when you call `op(...)`. ## Multi-Head Latent Attention -::: tileops.attention.MLADecodeWithKVCacheFwdOp +::: tileops.attention.MLADenseDecodeFwdOp options: show_root_heading: true heading_level: 3 @@ -93,7 +93,7 @@ what runs when you call `op(...)`. ## DeepSeek Sparse Attention -::: tileops.attention.DSADecodeWithKVCacheFwdOp +::: tileops.attention.DSADenseFwdOp options: show_root_heading: true heading_level: 3 diff --git a/docs/api/elementwise.md b/docs/api/elementwise.md index bf95f5ea..66fe078c 100644 --- a/docs/api/elementwise.md +++ b/docs/api/elementwise.md @@ -179,37 +179,37 @@ what runs when you call `op(...)`. ## Activations -::: tileops.elementwise.ReluFwdOp +::: tileops.elementwise.ReLUFwdOp options: show_root_heading: true heading_level: 3 members: ["__init__", "forward"] -::: tileops.elementwise.LeakyReluFwdOp +::: tileops.elementwise.LeakyReLUFwdOp options: show_root_heading: true heading_level: 3 members: ["__init__", "forward"] -::: tileops.elementwise.EluFwdOp +::: tileops.elementwise.ELUFwdOp options: show_root_heading: true heading_level: 3 members: ["__init__", "forward"] -::: tileops.elementwise.SeluFwdOp +::: tileops.elementwise.SELUFwdOp options: show_root_heading: true heading_level: 3 members: ["__init__", "forward"] -::: tileops.elementwise.GeluFwdOp +::: tileops.elementwise.GELUFwdOp options: show_root_heading: true heading_level: 3 members: ["__init__", "forward"] -::: tileops.elementwise.SiluFwdOp +::: tileops.elementwise.SiLUFwdOp options: show_root_heading: true heading_level: 3 @@ -257,7 +257,7 @@ what runs when you call `op(...)`. heading_level: 3 members: ["__init__", "forward"] -::: tileops.elementwise.PreluFwdOp +::: tileops.elementwise.PReLUFwdOp options: show_root_heading: true heading_level: 3 @@ -271,19 +271,19 @@ what runs when you call `op(...)`. ## Gated activations -::: tileops.elementwise.SiluAndMulFwdOp +::: tileops.elementwise.SiLUAndMulFwdOp options: show_root_heading: true heading_level: 3 members: ["__init__", "forward"] -::: tileops.elementwise.GeluAndMulFwdOp +::: tileops.elementwise.GELUAndMulFwdOp options: show_root_heading: true heading_level: 3 members: ["__init__", "forward"] -::: tileops.elementwise.GeluTanhAndMulFwdOp +::: tileops.elementwise.GELUTanhAndMulFwdOp options: show_root_heading: true heading_level: 3 diff --git a/docs/api/index.md b/docs/api/index.md index 677b8b7e..374bd8e0 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -7,9 +7,9 @@ dtypes — and the call takes the tensors. Both appear under each op as `__init_ ```python import torch -from tileops.gemm import GemmFwdOp +from tileops.gemm import GEMMFwdOp -op = GemmFwdOp() # construct once, reuse +op = GEMMFwdOp() # construct once, reuse d = op(a, b) # the specialized kernel is built on first call ``` @@ -51,4 +51,4 @@ Two things this reference does not carry: These pages are generated from the docstrings in TileOPs, so an op whose docstring is thin reads thin here. The fix belongs upstream, in -[`src/tileops/ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/ops). +[`src/tileops/_ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/_ops). diff --git a/docs/api/linear-algebra.md b/docs/api/linear-algebra.md index d84e4674..ec3f9c49 100644 --- a/docs/api/linear-algebra.md +++ b/docs/api/linear-algebra.md @@ -7,7 +7,7 @@ what runs when you call `op(...)`. ## Dense GEMM -::: tileops.gemm.GemmFwdOp +::: tileops.gemm.GEMMFwdOp options: show_root_heading: true heading_level: 3 @@ -15,7 +15,7 @@ what runs when you call `op(...)`. merge_init_into_class: false show_signature_annotations: false -::: tileops.gemm.GemmFP8FwdOp +::: tileops.gemm.GEMMFP8FwdOp options: show_root_heading: true heading_level: 3 @@ -23,7 +23,7 @@ what runs when you call `op(...)`. merge_init_into_class: false show_signature_annotations: false -::: tileops.gemm.GemmW4A16FwdOp +::: tileops.gemm.GEMMW4A16FwdOp options: show_root_heading: true heading_level: 3 @@ -31,7 +31,7 @@ what runs when you call `op(...)`. merge_init_into_class: false show_signature_annotations: false -::: tileops.gemm.GemmW4A8FwdOp +::: tileops.gemm.GEMMW4A8FwdOp options: show_root_heading: true heading_level: 3 @@ -39,7 +39,7 @@ what runs when you call `op(...)`. merge_init_into_class: false show_signature_annotations: false -::: tileops.gemm.GemmINT8W8A8FwdOp +::: tileops.gemm.GEMMINT8W8A8FwdOp options: show_root_heading: true heading_level: 3 @@ -49,7 +49,7 @@ what runs when you call `op(...)`. ## Batched GEMM -::: tileops.gemm.BmmFwdOp +::: tileops.gemm.BMMFwdOp options: show_root_heading: true heading_level: 3 @@ -57,7 +57,7 @@ what runs when you call `op(...)`. merge_init_into_class: false show_signature_annotations: false -::: tileops.gemm.BmmFP8FwdOp +::: tileops.gemm.BMMFP8FwdOp options: show_root_heading: true heading_level: 3 @@ -67,7 +67,7 @@ what runs when you call `op(...)`. ## Grouped GEMM -::: tileops.gemm.GroupedGemmFwdOp +::: tileops.gemm.GroupedGEMMFwdOp options: show_root_heading: true heading_level: 3 diff --git a/docs/api/moe.md b/docs/api/moe.md index 414a055a..766a3d49 100644 --- a/docs/api/moe.md +++ b/docs/api/moe.md @@ -8,7 +8,7 @@ what runs when you call `op(...)`. A routed mixture-of-experts layer is available two ways here. `FusedMoEFwdOp` runs the whole FFN, and `FusedMoESharedExpertFwdOp` adds a shared expert beside the routed ones. The rest are its stages, callable on their own: the op that picks each token's experts, the ops that move tokens into an expert-contiguous layout and -back, and the expert GEMMs that run on it. `MoEGroupedGemmFwdOp` is one grouped +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, and `IndexedExpertMLPFwdOp` is the backend it picks instead when @@ -56,13 +56,13 @@ the routes are few enough to read the weights once per route rather than once pe ## Expert GEMMs -::: tileops.moe.MoEGroupedGemmFwdOp +::: tileops.moe.MoEGroupedGEMMFwdOp options: show_root_heading: true heading_level: 3 members: ["__init__", "forward"] -::: tileops.moe.MoEGroupedGemmFP8FwdOp +::: tileops.moe.MoEGroupedGEMMFP8FwdOp options: show_root_heading: true heading_level: 3 diff --git a/docs/api/quantization.md b/docs/api/quantization.md index 5a869add..2f8ef130 100644 --- a/docs/api/quantization.md +++ b/docs/api/quantization.md @@ -19,7 +19,7 @@ what runs when you call `op(...)`. heading_level: 3 members: ["__init__", "forward"] -::: tileops.quantization.INT8QuantPerBlockFwdOp +::: tileops.quantization.INT8QuantPerTokenGroupFwdOp options: show_root_heading: true heading_level: 3 @@ -33,7 +33,7 @@ what runs when you call `op(...)`. ## FP8 quantization -::: tileops.quantization.FP8QuantFwdOp +::: tileops.quantization.FP8QuantPerTokenGroupFwdOp options: show_root_heading: true heading_level: 3 diff --git a/docs/api/trace.md b/docs/api/trace.md index a95b8c60..e812e571 100644 --- a/docs/api/trace.md +++ b/docs/api/trace.md @@ -6,7 +6,7 @@ wire the markers into a `@tilelang.jit` builder, run, and export a timeline. For a walkthrough, see the [In-Kernel Timeline Trace guide](../performance-guides/trace-timeline.md). -::: tileops.trace.api._Trace +::: tileops.trace._api._Trace options: show_root_heading: false show_root_toc_entry: false diff --git a/docs/backends.md b/docs/backends.md index 72f043b4..3f41ab59 100644 --- a/docs/backends.md +++ b/docs/backends.md @@ -31,20 +31,13 @@ The second half explains why the protocol is designed this way: - the interfaces available to a caller; - what the protocol deliberately does not support. -## Two ways to extend dispatch {#two-ways} +## How a backend extends dispatch {#extend} -A package outside TileOPs picks one of two mechanisms, by how much of an op it takes -over: - -1. `register_kernel_type`: adds an implementation to a kernel interface, which takes - part in selection alongside the in-tree implementations; -1. target: takes over every call of the op. - -In the first, the kernel class follows the kernel interface, the same contract the -in-tree implementations follow; how to write one is in -[How a backend joins TileOPs](user-guide/dispatch/backends.md). A target follows the op's -manifest signature instead, and the backend writes a `build_kernel` for it. The rest of -this page covers the target only. +A package outside TileOPs takes over every call of an op through a target. A target +follows the op's manifest signature, and the backend writes a `build_kernel` for it. The +kernel interfaces and call specs of TileOPs' ops are not public API, so a backend does not +add an implementation to a single kernel interface; see +[How a backend joins TileOPs](user-guide/dispatch/backends.md). ## Four things to write @@ -126,10 +119,10 @@ this call arrives as `None`. ```python # GroupNormFwdOp's spec: weight and bias are optional, and arrive as None when absent -def build_group_norm(x, weight, bias, *, num_groups, eps): +def build_group_norm(input, weight, bias, *, num_groups, eps): if weight is None: # presence read off the slot - return AcmeGroupNorm(num_groups, eps, x.dtype) - return AcmeGroupNormAffine(num_groups, eps, x.dtype) + return AcmeGroupNorm(num_groups, eps, input.dtype) + return AcmeGroupNormAffine(num_groups, eps, input.dtype) ``` The op layer calls `build_kernel` with the call's device made current through @@ -187,8 +180,8 @@ call's device: `arch`, `calibration`, `sm_count` and `smem_budget`. | Another registered type | The registered `read`, which takes the call's `torch.device` and returns `DeviceFacts` | | An unregistered type, such as `cpu` | Nowhere: it reports `-1 / None / 0 / 0` | -- A backend registers facts when it adds implementations with `register_kernel_type` that - are selected by architecture, SM count or shared memory. An implementation that declares +- Facts are read by implementations that are selected by architecture, SM count or shared + memory. An implementation that declares `supported_archs` is never available on an unregistered device type. - A target that serves whole ops does not need them: its kernels make their own choices. - Facts belong to a device type, not to a target. If several targets serve one device @@ -266,7 +259,7 @@ The steps of one call, from user code to a backend's `build_kernel`: ```python # ── the caller ─────────────────────────────────────────────────────── -op = GemmFwdOp() # no target= in the constructor, so the inputs' device decides +op = GEMMFwdOp() # no target= in the constructor, so the inputs' device decides # target="acme" skips detection and uses it directly; # target=BUILTIN forces the in-tree kernels a = torch.randn(4096, 4096, dtype=torch.float16, device="acme:0") @@ -282,7 +275,7 @@ 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: run the checks generated from the manifest signature, then hand the whole op to the target ── -# GemmFwdOp's own forward and kernel_for serve the in-tree path only, and do not run +# GEMMFwdOp's own forward and kernel_for serve the in-tree path only, and do not run # the tensors go to the target in signature.inputs order, inputs it does not write made contiguous # ── op layer: look up the external memo table — device, then input signature ── @@ -311,7 +304,7 @@ def build_gemm(a: TensorSpec, b: TensorSpec, *, trans_a, trans_b): return AcmeGemm(a, b, trans_a, trans_b) -register_kernel_builder(op="GemmFwdOp", target="acme", build_kernel=build_gemm) +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 @@ -409,7 +402,7 @@ every op the target takes over, with each key spelled as in the manifest. The pa ```python BUILDERS = { "RMSNormFwdOp": build_rms_norm, - "GemmFwdOp": build_gemm, + "GEMMFwdOp": build_gemm, } ``` @@ -658,8 +651,7 @@ print(load_failures()) from its own set.** The second layer happens inside `build_kernel`, without the protocol: this path has 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 smaller mechanism, `register_kernel_type` (see -[Two ways to extend dispatch](#two-ways)); a target bypasses it. +precedence) belongs to the in-tree path only; a target bypasses it. `detect` answers only which devices belong to the backend, and nothing finer. **Whether this call is supported (dtype, shape, parameter combination) is answered by @@ -762,7 +754,7 @@ A claim comes with three sets of numbers: `pytest tests/ops --collect-only -q` runs: one with the default options, one with `--tileops-target` and `--tileops-device`. -The in-tree kernels ([`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels)) +The in-tree kernels ([`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 in-tree implementations run by default.** With no backend installed, no `target=` diff --git a/docs/backends.zh.md b/docs/backends.zh.md index 4d4bf983..8b470610 100644 --- a/docs/backends.zh.md +++ b/docs/backends.zh.md @@ -25,14 +25,9 @@ TileLang 是支持多种 backend 的 DSL，每种硬件各有一套独立的 ker - 调用方可用的接口； - 刻意不支持的情形。 -## 两种接入方式 {#two-ways} +## 接入方式 {#extend} -一个仓库之外的包按自己接管的范围，从两种方式中选一种： - -1. `register_kernel_type`：为一个 kernel 接口新增一个实现，与 in-tree 实现一起参与选择； -2. target：接管 op 的全部调用。 - -第一种方式写出的 kernel 类遵守 kernel 接口，与 in-tree 实现遵守同一份契约，写法见 [backend 如何接入](user-guide/dispatch/backends.md)。target 遵守的是 op 在 manifest 中的签名，backend 为它写一个 `build_kernel`。本页以下各节只讲 target。 +一个仓库之外的包通过 target 接管 op 的全部调用。target 遵守的是 op 在 manifest 中的签名，backend 为它写一个 `build_kernel`。TileOPs 各 op 的 kernel 接口与 call spec 不是公开 API，backend 不为单个 kernel 接口新增实现，见 [backend 如何接入](user-guide/dispatch/backends.md)。 ## 写一个 backend 要做的四件事 @@ -99,10 +94,10 @@ def build_kernel(*inputs: "TensorSpec | None", **params) -> Callable[..., Kernel ```python # GroupNormFwdOp 的 spec：weight、bias 是可选输入，没传时实参是 None -def build_group_norm(x, weight, bias, *, num_groups, eps): +def build_group_norm(input, weight, bias, *, num_groups, eps): if weight is None: # 由该位置的实参判断是否传入 - return AcmeGroupNorm(num_groups, eps, x.dtype) - return AcmeGroupNormAffine(num_groups, eps, x.dtype) + return AcmeGroupNorm(num_groups, eps, input.dtype) + return AcmeGroupNormAffine(num_groups, eps, input.dtype) ``` op 层调用 `build_kernel` 时，经 `torch.accelerator` 把调用设备设为当前设备；调用设备不是进程 accelerator 的类型时不切换。 @@ -151,7 +146,7 @@ def register_device_facts(device_type: str, read: Callable[[torch.device], Devic | 已注册的其他类型 | 注册的 `read`：收到调用的 `torch.device`，返回 `DeviceFacts` | | 未注册的类型，如 `cpu` | 没有设备事实：`-1 / None / 0 / 0` | -- 需要注册的是用 `register_kernel_type` 加入、按架构、SM 数或共享内存选择的实现。声明了 `supported_archs` 的实现在未注册的设备类型上永远不可用。 +- 读取设备事实的是按架构、SM 数或共享内存选择的实现。声明了 `supported_archs` 的实现在未注册的设备类型上永远不可用。 - 只服务整个 op 的 target 不需要注册：target 的 kernel 自己决定怎么选。 - 按设备类型注册，与 target 无关：同一种设备上有多个 target 时，只由一个包注册。重复注册、注册 `cuda` 都报 `BackendError`。 - `arch` 只在同一种设备类型内可比，所以声明 `supported_archs` 的实现，其 `devices` 只能有一种设备类型。 @@ -213,7 +208,7 @@ def build_gemm(a: TensorSpec, b: TensorSpec, *, trans_a, trans_b): ```python # ── 调用方 ─────────────────────────────────────────────────────────── -op = GemmFwdOp() # 构造时未给出 target=，由输入张量的设备决定 +op = GEMMFwdOp() # 构造时未给出 target=，由输入张量的设备决定 # 给出 target="acme" 时跳过设备探测，直接使用它； # 给出 target=BUILTIN 时强制使用 in-tree kernel a = torch.randn(4096, 4096, dtype=torch.float16, device="acme:0") @@ -229,7 +224,7 @@ d = op(a, b) # 所有输入必须在同一设备上：a.devi # 两个以上返回 True → 抛出 AmbiguousTargetError，要求显式给出 target= # ── op 层：先执行由 manifest 签名生成的检查，再由 target 服务整个 op ── -# GemmFwdOp 自己的 forward 与 kernel_for 只服务 in-tree 实现，这次不执行 +# GEMMFwdOp 自己的 forward 与 kernel_for 只服务 in-tree 实现，这次不执行 # 传给 target 的张量按 signature.inputs 排列，不写入的输入先转为连续张量 # ── op 层：按设备与输入签名查外部记忆表 ───────────────────────────── @@ -258,7 +253,7 @@ def build_gemm(a: TensorSpec, b: TensorSpec, *, trans_a, trans_b): return AcmeGemm(a, b, trans_a, trans_b) -register_kernel_builder(op="GemmFwdOp", target="acme", build_kernel=build_gemm) +register_kernel_builder(op="GEMMFwdOp", target="acme", build_kernel=build_gemm) ``` `build_gemm` 由 op 层调用，backend 自己不调用它。import backend 模块时只把它登记进注册表；一次调用由这个 target 服务且外部记忆表未命中时，op 层才调用它，每个「设备 + 输入签名」一次。它返回的可调用对象由 op 层 launch 并存入记忆表。 @@ -325,7 +320,7 @@ def build_rms_norm(x: TensorSpec, weight: TensorSpec | None, *, normalized_shape ```python BUILDERS = { "RMSNormFwdOp": build_rms_norm, - "GemmFwdOp": build_gemm, + "GEMMFwdOp": build_gemm, } ``` @@ -528,7 +523,7 @@ print(load_failures()) | **`detect`** | backend 编写的一个函数，每个 target 一个 | 第一层的选择依据：接收一个 `torch.device`，回答这类设备是否是自己这套 kernel 的目标设备；不是则返回 `False` | | **`build_kernel`** | backend 为某个 op 编写的一个函数，每组 `(op, target)` 一个 | 第二层：接收本次调用的描述，即各输入张量的 device、dtype、shape 与 op 参数，从自己这套 kernel 中选定一个，构造好并返回 | -**选择分两层：TileOPs 选 target，target 在自己那套 kernel 中选一个。** 第二层发生在 `build_kernel` 内部，协议不参与，这条路径上没有 kernel 一级的概念、能力协商与候选筛选。TileOPs 实际执行的候选筛选（可用性、适用范围、优先关系）属于 in-tree 实现与范围较小的 `register_kernel_type`（见[两种接入方式](#two-ways)），target 绕过这些筛选。 +**选择分两层：TileOPs 选 target，target 在自己那套 kernel 中选一个。** 第二层发生在 `build_kernel` 内部，协议不参与，这条路径上没有 kernel 一级的概念、能力协商与候选筛选。TileOPs 实际执行的候选筛选（可用性、适用范围、优先关系）只属于 in-tree 实现，target 绕过这些筛选。 `detect` 只回答设备的归属，不回答更细的问题。**本次调用是否受支持（涉及 dtype、形状与参数组合）由 `build_kernel` 回答**，因为只有它看得到完整的输入描述与参数，不支持时也在那里报错。`detect` 只拿到一个 `torch.device`，无法作出这些判断。 @@ -613,7 +608,7 @@ backend 对一个 op 的结论是「通过了这台设备上实际运行的测 - 各测试的 `checked_outputs`； - 按测试文件统计的被去掉的测试数。被去掉的测试不进 JUnit，这个数由 backend 自行计算，即两次 `pytest tests/ops --collect-only -q` 的 node id 之差：一次用默认参数，一次加上 `--tileops-target` 与 `--tileops-device`。 -TileOPs 的 in-tree kernel（[`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels)）是**默认实现**：它没有 target 名，也不进入注册表。 +TileOPs 的 in-tree kernel（[`src/tileops/_kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/_kernels)）是**默认实现**：它没有 target 名，也不进入注册表。 **默认使用 in-tree 实现。** 没有安装 backend、没有通过 `target=` 指名、也没有设置进程默认值时，调用由 in-tree 实现服务。安装一个 backend 之后，只有当它的 `detect` 认领了这块设备，或者它被 `target=` 或 `set_default_target` 指名时，该 op 的 kernel 才换成这个 backend 的那一套。 diff --git a/docs/design/architecture.md b/docs/design/architecture.md deleted file mode 100644 index 6c473840..00000000 --- a/docs/design/architecture.md +++ /dev/null @@ -1,3 +0,0 @@ -{% - include-markdown "../../TileOPs/docs/design/architecture.md" -%} diff --git a/docs/design/index.md b/docs/design/index.md index 42e8e830..d958247b 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -6,8 +6,6 @@ repository — the source of truth — pulled in at site build time. ## System structure -- [Architecture](architecture.md) - Top-level module layout and the spec-driven pipeline. - [Layer Boundaries](layer-boundaries.md) What each layer owns and depends on, and the interfaces layers compose through. diff --git a/docs/index.md b/docs/index.md index 74686c85..db73d29f 100644 --- a/docs/index.md +++ b/docs/index.md @@ -32,12 +32,12 @@ to the call, and the specialized kernel is compiled and cached on the first call ```python import torch -from tileops.gemm import GemmFwdOp +from tileops.gemm import GEMMFwdOp a = torch.randn(4096, 4096, device="cuda", dtype=torch.float16) b = torch.randn(4096, 4096, device="cuda", dtype=torch.float16) -op = GemmFwdOp() # NT by default: a=[M, K], b=[N, K] +op = GEMMFwdOp() # NT by default: a=[M, K], b=[N, K] d = op(a, b) # -> [M, N] flops, nbytes = op.eval_roofline() # what the call had to do and move ``` diff --git a/docs/index.zh.md b/docs/index.zh.md index f0fe0df5..772fb317 100644 --- a/docs/index.zh.md +++ b/docs/index.zh.md @@ -18,12 +18,12 @@ op 在构造时不绑定任何形状。形状和 dtype 取自调用时传入的 ```python import torch -from tileops.gemm import GemmFwdOp +from tileops.gemm import GEMMFwdOp a = torch.randn(4096, 4096, device="cuda", dtype=torch.float16) b = torch.randn(4096, 4096, device="cuda", dtype=torch.float16) -op = GemmFwdOp() # 默认 NT 布局：a=[M, K], b=[N, K] +op = GEMMFwdOp() # 默认 NT 布局：a=[M, K], b=[N, K] d = op(a, b) # -> [M, N] flops, nbytes = op.eval_roofline() # 本次调用所需的计算量与访存量 ``` diff --git a/docs/new-op.md b/docs/new-op.md index 7858e23a..40a82d11 100644 --- a/docs/new-op.md +++ b/docs/new-op.md @@ -10,13 +10,13 @@ the other five places are written from it.**{ .keystone } | # | File | Held to the spec by | Contents | | --- | --- | --- | --- | | 1 | [`src/tileops/manifest/spec/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/manifest/spec)`.yaml` | the validator's `schema` and `signature` levels | the spec itself | -| 2 | [`src/tileops/ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/ops)`/…` | the validator, against `__init__`, `forward` and the declared kernels; the checks generated around every call | the op class, subclassing `Op` | -| 2 | [`src/tileops/ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/ops)`/__init__.py` and [`src/tileops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops)`.py` | the validator: the family's `__all__` agrees with the manifest | the op's name, exported by its family and on the public path `tileops..` | -| 3 | [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels)`/…` | — | the kernel classes, subclassing `Kernel` | +| 2 | [`src/tileops/_ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/_ops)`/…` | the validator, against `__init__`, `forward` and the declared kernels; the checks generated around every call | the op class, subclassing `Op` | +| 2 | [`src/tileops/_ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/_ops)`/__init__.py` and [`src/tileops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops)`.py` | the validator: the family's `__all__` agrees with the manifest | the op's name, exported by its family and on the public path `tileops..` | +| 3 | [`src/tileops/_kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/_kernels)`/…` | — | the kernel classes, subclassing `Kernel` | | 4 | [`tests/ops/`](https://github.com/tile-ai/TileOPs/tree/main/tests/ops)`test_.py` | the contract tests, which run every workload row | the numerical comparison against the reference `ref_program` | | 5 | [`benchmarks/ops/`](https://github.com/tile-ai/TileOPs/tree/main/benchmarks/ops)`bench_.py` | the validator's `bench` level | the benchmark, and the op's case entry in `benchmarks/_cases/` | -The steps below take `GemmFwdOp`, the simplest matmul, through all six places. +The steps below take `GEMMFwdOp`, the simplest matmul, through all six places. ## Step 1: write the spec @@ -26,10 +26,10 @@ What each field means and how to write it is in yet. The checks that read code are skipped at this status, so they do not fail on the missing class. -`GemmFwdOp`'s spec, with one workload row: +`GEMMFwdOp`'s spec, with one workload row: ```yaml -GemmFwdOp: +GEMMFwdOp: ref_api: torch.matmul family: gemm status: spec-only @@ -62,23 +62,25 @@ and is declared on the op class in step 2. ## Step 2: write the op class {#op-class} -The op class subclasses [`Op`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops/op_base.py) and sits between the spec and the kernel. The checks +The op class subclasses [`Op`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/_ops/op_base.py) and sits between the spec and the kernel. The checks around every call (dtypes, shapes, the refinements and output shape inference) are generated from the signature when the class is defined, and the class writes none of -them. What the class writes is how a call reaches a kernel. +them. What the class writes is how a call reaches a kernel. Code outside TileOPs that +subclasses `Op` imports `Op`, `Kernel`, `KernelInterface`, `CallSpec` and `Entry` from +`tileops.authoring`. ### The class, and its members -`GemmFwdOp`'s skeleton, with docstrings elided: +`GEMMFwdOp`'s skeleton, with docstrings elided: ```python -class GemmFwdOp(Op): +class GEMMFwdOp(Op): kernel_types: ClassVar[Mapping[str, type[Kernel]]] = { - "gemm_tma": GemmTMAKernel, - "gemm_cp_async": GemmCpAsyncKernel, + "gemm_tma": GEMMTMAKernel, + "gemm_cp_async": GEMMCpAsyncKernel, "gemv": GemvKernel, } - interfaces: ClassVar[Mapping[str, type[KernelInterface]]] = {"gemm": GemmFwdInterface} + interfaces: ClassVar[Mapping[str, type[KernelInterface]]] = {"gemm": GEMMFwdInterface} def __init__(self, trans_a=False, trans_b=True, *, target=None): self.trans_a = trans_a @@ -88,7 +90,7 @@ class GemmFwdOp(Op): def 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 - call = GemmCall( # 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, @@ -135,7 +137,7 @@ before is returned from the table, and otherwise the kernel is built and stored. `kernel_for` takes two arguments: - **`interface`**: a key of `interfaces`, naming one kernel call the op makes. - `GemmFwdOp` makes one kernel call, so it declares one interface, `"gemm"`. A second + `GEMMFwdOp` makes one kernel call, so it declares one interface, `"gemm"`. A second kernel interface is added 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 kernel that is faster on some shape range or some architecture @@ -150,7 +152,7 @@ The op calls the returned kernel with the parameters of the kernel interface's abstract `forward`, in that order. A kernel interface is a class in -[`src/tileops/kernels//call_spec.py`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels), +[`src/tileops/_kernels//call_spec.py`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/_kernels), next to the call spec it names in `request`; a family with one kernel file keeps both in that file. Its name has the form `{Name}{Fwd|Bwd}Interface`, with variant words before the direction, and `interface-names-lint` checks it. Its abstract `forward` is @@ -159,10 +161,10 @@ so its docstring states each tensor's shape, dtype, memory layout, device, and w it is written in place: ```python -class GemmFwdInterface(KernelInterface): +class GEMMFwdInterface(KernelInterface): """Dense matmul under the ``(trans_a, trans_b)`` layout the call states.""" - request = GemmCall + request = GEMMCall @abstractmethod def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: @@ -199,7 +201,7 @@ op is in [How a backend joins TileOPs](user-guide/dispatch/backends.md). Add the op's name to the imports and `__all__` in two places: -1. its family's [`src/tileops/ops//__init__.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops), where the class is implemented; +1. its family's [`src/tileops/_ops//__init__.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/_ops), where the class is implemented; 1. [`src/tileops/.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops), the public path. Without the second, `from tileops. import ...` does not find the op, and the @@ -207,7 +209,7 @@ API reference does not include 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) and the kernel interface it implements, 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 kernel interface it implements, lives under [`src/tileops/_kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/_kernels), is written in TileLang. Construction obtains the TileLang builder; the program compiles the first time a call runs it with a given config. It implements `forward`, which the base class's `__call__` runs. Its constructor is called by the builder its own `entry_for` @@ -218,11 +220,11 @@ The kernel is the only one of the six places the spec does not constrain: a kern neither reads the spec nor is checked against it. The split of arguments between the constructor and the call is a hard requirement: -**only values compiled into the generated code go in the constructor.** `GemmTMAKernel` +**only values compiled into the generated code go in the constructor.** `GEMMTMAKernel` splits them like this: ```python -class GemmTMAKernel(Kernel, GemmFwdInterface): +class GEMMTMAKernel(Kernel, GEMMFwdInterface): def __init__(self, m, n, k, dtype, config=None, trans_a=False, trans_b=False, ...): self.kernel = _gemm_kernel(m, n, k, trans_a, trans_b, self.dtype_str, ...) # the TileLang builder self.init_config(config) # tile sizes and pipeline depth @@ -291,12 +293,12 @@ file calls both: import pytest from benchmarks import api as bench -from tileops.ops import GemmFwdOp +from tileops.gemm import GEMMFwdOp -@pytest.mark.parametrize("case", bench.cases(GemmFwdOp), ids=lambda case: case.id) +@pytest.mark.parametrize("case", bench.cases(GEMMFwdOp), ids=lambda case: case.id) def test_gemm_bench(case) -> None: - op = GemmFwdOp(**case.arguments) + op = GEMMFwdOp(**case.arguments) bench.Runner(op, case).compare({"torch-cublas": case.reference}) ``` @@ -317,7 +319,7 @@ is in [How a benchmark is timed](timing.md). With the other five places written, run the three commands below to check the work: ```bash -python scripts/validate_manifest.py --check-op GemmFwdOp # spec and code agree +python scripts/validate_manifest.py --check-op GEMMFwdOp # spec and code agree python -m pytest tests/ops/test_gemm.py -v # numerics match ref_program python -m pytest benchmarks/ops/bench_gemm.py # the benchmark produces numbers ``` diff --git a/docs/new-op.zh.md b/docs/new-op.zh.md index 7e90b56d..9de1588d 100644 --- a/docs/new-op.zh.md +++ b/docs/new-op.zh.md @@ -7,22 +7,22 @@ spec 最先编写，因为后面五个文件的内容都由 spec 决定，最后 | # | 文件 | 由谁对照 spec 检查 | 内容 | | --- | --- | --- | --- | | 1 | [`src/tileops/manifest/spec/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/manifest/spec)`.yaml` | 校验器的 `schema` 与 `signature` 两级 | spec 本身 | -| 2 | [`src/tileops/ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/ops)`/…` | 校验器对照 `__init__`、`forward` 与声明的 kernel；每次调用前后生成的检查 | op 类，继承 `Op` | -| 2 | [`src/tileops/ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/ops)`/__init__.py` 与 [`src/tileops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops)`.py` | 校验器：family 的 `__all__` 与 manifest 一致 | op 名，由所属 family 导出，并出现在公开路径 `tileops..` 上 | -| 3 | [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels)`/…` | —— | kernel 类，继承 `Kernel` | +| 2 | [`src/tileops/_ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/_ops)`/…` | 校验器对照 `__init__`、`forward` 与声明的 kernel；每次调用前后生成的检查 | op 类，继承 `Op` | +| 2 | [`src/tileops/_ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/_ops)`/__init__.py` 与 [`src/tileops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops)`.py` | 校验器：family 的 `__all__` 与 manifest 一致 | op 名，由所属 family 导出，并出现在公开路径 `tileops..` 上 | +| 3 | [`src/tileops/_kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/_kernels)`/…` | —— | kernel 类，继承 `Kernel` | | 4 | [`tests/ops/`](https://github.com/tile-ai/TileOPs/tree/main/tests/ops)`test_<名字>.py` | 契约测试，逐个运行每个 workload 行 | 与参考实现 `ref_program` 的数值比对 | | 5 | [`benchmarks/ops/`](https://github.com/tile-ai/TileOPs/tree/main/benchmarks/ops)`bench_<名字>.py` | 校验器的 `bench` 级 | benchmark，以及该 op 在 `benchmarks/_cases/` 中的 case 注册项 | -下文以最简单的矩阵乘 `GemmFwdOp` 为例，依次说明这六处。 +下文以最简单的矩阵乘 `GEMMFwdOp` 为例，依次说明这六处。 ## 第一步：写 spec spec 各字段的含义与写法见[读写 manifest](user-guide/manifest/index.md)。新 op 的 spec 先写成 `status: spec-only`，表示接口已经确定、实现尚未完成。此时所有需要读取代码的检查都会跳过，不会因为找不到类而报错。 -`GemmFwdOp` 的 spec 只列出一行 workload： +`GEMMFwdOp` 的 spec 只列出一行 workload： ```yaml -GemmFwdOp: +GEMMFwdOp: ref_api: torch.matmul family: gemm status: spec-only @@ -54,20 +54,20 @@ spec 不写文件路径，也不写 kernel。由哪些 kernel 服务这个 op ## 第二步：写 op 类 {#op-class} -op 类继承 [`Op`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops/op_base.py)，位于 spec 与 kernel 之间。每次调用前后的检查，包括 dtype、形状、约束与输出形状推导，都在类定义时依照签名生成，op 类不手写其中任何一项。op 类需要编写的是一次调用如何到达 kernel。 +op 类继承 [`Op`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/_ops/op_base.py)，位于 spec 与 kernel 之间。每次调用前后的检查，包括 dtype、形状、约束与输出形状推导，都在类定义时依照签名生成，op 类不手写其中任何一项。op 类需要编写的是一次调用如何到达 kernel。TileOPs 之外继承 `Op` 的代码从 `tileops.authoring` 导入 `Op`、`Kernel`、`KernelInterface`、`CallSpec` 与 `Entry`。 ### 类的骨架与成员 -`GemmFwdOp` 的骨架如下，略去 docstring： +`GEMMFwdOp` 的骨架如下，略去 docstring： ```python -class GemmFwdOp(Op): +class GEMMFwdOp(Op): kernel_types: ClassVar[Mapping[str, type[Kernel]]] = { - "gemm_tma": GemmTMAKernel, - "gemm_cp_async": GemmCpAsyncKernel, + "gemm_tma": GEMMTMAKernel, + "gemm_cp_async": GEMMCpAsyncKernel, "gemv": GemvKernel, } - interfaces: ClassVar[Mapping[str, type[KernelInterface]]] = {"gemm": GemmFwdInterface} + interfaces: ClassVar[Mapping[str, type[KernelInterface]]] = {"gemm": GEMMFwdInterface} def __init__(self, trans_a=False, trans_b=True, *, target=None): self.trans_a = trans_a @@ -77,7 +77,7 @@ class GemmFwdOp(Op): def 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 - call = GemmCall( # 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, @@ -112,18 +112,18 @@ kernel 是编译产物，构造一次需要几百毫秒到几秒，而一个 op `kernel_for` 接受两个参数： -- **`interface`**：`interfaces` 的一个 key，指 op 发出的某一个 kernel 调用。`GemmFwdOp` 只发出一个 kernel 调用，因此只声明一个 `"gemm"`。只有语义或调用契约改变时才新增一个 kernel 接口，例如 `BatchNormFwdOp` 的 `batch_norm_fwd_train` 与 `batch_norm_fwd_infer` 返回的内容不同。在某个形状范围或某个架构上更快的 kernel 是已有 kernel 接口的另一个实现。 +- **`interface`**：`interfaces` 的一个 key，指 op 发出的某一个 kernel 调用。`GEMMFwdOp` 只发出一个 kernel 调用，因此只声明一个 `"gemm"`。只有语义或调用契约改变时才新增一个 kernel 接口，例如 `BatchNormFwdOp` 的 `batch_norm_fwd_train` 与 `batch_norm_fwd_infer` 返回的内容不同。在某个形状范围或某个架构上更快的 kernel 是已有 kernel 接口的另一个实现。 - **`call`**：一个冻结的 `CallSpec` 子类，包含选择实现和构造 kernel 所需的信息，即形状、dtype、op 的语义参数与设备。它必须是这个 kernel 接口的 `request` 指定的类型。设备事实（`arch`、`sm_count`、`calibration`、`smem_budget`）由派发机制在缓存未命中时从 `call.device` 推出，调用方不填写。 op 按 kernel 接口的抽象 `forward` 方法所声明的参数及其顺序调用取得的 kernel。 -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` 检查。kernel 接口的抽象 `forward` 是所有实现（in-tree 实现与 backend 提供的实现）共同依据的唯一契约，因此它的 docstring 写明每个张量的形状、dtype、内存布局、设备，以及是否被原地写入： +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` 检查。kernel 接口的抽象 `forward` 是所有实现（in-tree 实现与 backend 提供的实现）共同依据的唯一契约，因此它的 docstring 写明每个张量的形状、dtype、内存布局、设备，以及是否被原地写入： ```python -class GemmFwdInterface(KernelInterface): +class GEMMFwdInterface(KernelInterface): """Dense matmul under the ``(trans_a, trans_b)`` layout the call states.""" - request = GemmCall + request = GEMMCall @abstractmethod def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: @@ -149,21 +149,21 @@ op 不编写 `entry_for`，也不自行维护 kernel 缓存。完全没有 in-tr op 名需要加入两处的导入与 `__all__`： -1. 所属 family 的 [`src/tileops/ops//__init__.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops)，即类的实现位置； +1. 所属 family 的 [`src/tileops/_ops//__init__.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/_ops)，即类的实现位置； 1. [`src/tileops/.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops)，即公开路径。 缺少后者时，`from tileops. import ...` 无法导入这个 op，API 参考也不会收录它。 ## 第三步：写 kernel -kernel 类继承 [`Kernel`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/kernel_base.py) 与它实现的 kernel 接口，放在 [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 下，用 TileLang 编写：构造时取得 TileLang 构建函数，程序在首次以某个配置调用时编译。kernel 类实现 `forward`，由基类的 `__call__` 调用。构造函数由本类 `entry_for` 返回的 builder 调用；`forward` 接受 kernel 接口规定的参数，也就是第二步中的 `kernel(a, b)`。 +kernel 类继承 [`Kernel`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/_kernels/kernel_base.py) 与它实现的 kernel 接口，放在 [`src/tileops/_kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/_kernels) 下，用 TileLang 编写：构造时取得 TileLang 构建函数，程序在首次以某个配置调用时编译。kernel 类实现 `forward`，由基类的 `__call__` 调用。构造函数由本类 `entry_for` 返回的 builder 调用；`forward` 接受 kernel 接口规定的参数，也就是第二步中的 `kernel(a, b)`。 kernel 是这六处中唯一不受 spec 约束的一处：kernel 不读取 spec，也不对照 spec 检查。 -构造参数与调用参数的划分有一条硬性要求：**只有会被编译进生成代码的值才作为构造参数。** `GemmTMAKernel` 的划分如下： +构造参数与调用参数的划分有一条硬性要求：**只有会被编译进生成代码的值才作为构造参数。** `GEMMTMAKernel` 的划分如下： ```python -class GemmTMAKernel(Kernel, GemmFwdInterface): +class GEMMTMAKernel(Kernel, GEMMFwdInterface): def __init__(self, m, n, k, dtype, config=None, trans_a=False, trans_b=False, ...): self.kernel = _gemm_kernel(m, n, k, trans_a, trans_b, self.dtype_str, ...) # the TileLang builder self.init_config(config) # tile sizes and pipeline depth @@ -211,12 +211,12 @@ benchmark 放在 [`benchmarks/ops/`](https://github.com/tile-ai/TileOPs/tree/mai import pytest from benchmarks import api as bench -from tileops.ops import GemmFwdOp +from tileops.gemm import GEMMFwdOp -@pytest.mark.parametrize("case", bench.cases(GemmFwdOp), ids=lambda case: case.id) +@pytest.mark.parametrize("case", bench.cases(GEMMFwdOp), ids=lambda case: case.id) def test_gemm_bench(case) -> None: - op = GemmFwdOp(**case.arguments) + op = GEMMFwdOp(**case.arguments) bench.Runner(op, case).compare({"torch-cublas": case.reference}) ``` @@ -229,7 +229,7 @@ benchmark 至少还需要记录一个非 TileOPs 的实现，否则这一行没 上述五处都写完之后，先运行以下三条命令自查： ```bash -python scripts/validate_manifest.py --check-op GemmFwdOp # spec and code agree +python scripts/validate_manifest.py --check-op GEMMFwdOp # spec and code agree python -m pytest tests/ops/test_gemm.py -v # numerics match ref_program python -m pytest benchmarks/ops/bench_gemm.py # the benchmark produces numbers ``` diff --git a/docs/performance-guides/memory-bound/index.md b/docs/performance-guides/memory-bound/index.md index 3756bb66..b63cc333 100644 --- a/docs/performance-guides/memory-bound/index.md +++ b/docs/performance-guides/memory-bound/index.md @@ -4,12 +4,12 @@ On a GPU, kernel speed is determined by whether compute or bandwidth becomes the bottleneck first. TileOPs measures a -**[calibration factor](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/perf/profiles/h200.yaml)** +**[calibration factor](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/_perf/profiles/h200.yaml)** with a [macro benchmark](https://github.com/tile-ai/TileOPs/tree/main/benchmarks/hardware): the hardware peak from the specification, multiplied by that factor, gives the effective peak a kernel can actually reach. That effective peak is the reference point for tuning. On an H200, the measured values are -[**57.27** TFLOP/s for fp32 FMA and **4.50** TB/s of memory bandwidth](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/perf/profiles/h200.yaml). +[**57.27** TFLOP/s for fp32 FMA and **4.50** TB/s of memory bandwidth](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/_perf/profiles/h200.yaml). The **ridge point** of the roofline is where the bandwidth slope meets the compute ceiling. Dividing compute by bandwidth places it at an arithmetic intensity of **12.72 flop/byte**, the number of floating-point operations per diff --git a/docs/performance-guides/memory-bound/index.zh.md b/docs/performance-guides/memory-bound/index.zh.md index 1ad1e589..b0ea5272 100644 --- a/docs/performance-guides/memory-bound/index.zh.md +++ b/docs/performance-guides/memory-bound/index.zh.md @@ -2,7 +2,7 @@ ## 什么是访存受限 -在 GPU 上，一个 kernel 的速度取决于算力与带宽哪一个先成为瓶颈。TileOPs 用 [macro benchmark](https://github.com/tile-ai/TileOPs/tree/main/benchmarks/hardware) 测算出一个**[校准系数](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/perf/profiles/h200.yaml)**：硬件 spec 给出的理论峰值乘以校准系数，得到实际可达的有效值，作为性能优化的参照。在 H200 上实测得到的 [fp32 FMA 的算力为 **57.27** TFLOP/s，访存带宽为 **4.50** TB/s](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/perf/profiles/h200.yaml)。roofline 的**拐点**（ridge point）是带宽斜线与算力上限这两段的交点，算力除以带宽得到它的横坐标，即拐点处的算术强度 **12.72 flop/byte**，也就是算力与带宽同时用满时每搬运一个字节对应的浮点运算次数。H200 的 roofline 如下图所示： +在 GPU 上，一个 kernel 的速度取决于算力与带宽哪一个先成为瓶颈。TileOPs 用 [macro benchmark](https://github.com/tile-ai/TileOPs/tree/main/benchmarks/hardware) 测算出一个**[校准系数](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/_perf/profiles/h200.yaml)**：硬件 spec 给出的理论峰值乘以校准系数，得到实际可达的有效值，作为性能优化的参照。在 H200 上实测得到的 [fp32 FMA 的算力为 **57.27** TFLOP/s，访存带宽为 **4.50** TB/s](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/_perf/profiles/h200.yaml)。roofline 的**拐点**（ridge point）是带宽斜线与算力上限这两段的交点，算力除以带宽得到它的横坐标，即拐点处的算术强度 **12.72 flop/byte**，也就是算力与带宽同时用满时每搬运一个字节对应的浮点运算次数。H200 的 roofline 如下图所示：
diff --git a/docs/performance-guides/trace-timeline.zh.md b/docs/performance-guides/trace-timeline.zh.md index 827edc38..9429e28a 100644 --- a/docs/performance-guides/trace-timeline.zh.md +++ b/docs/performance-guides/trace-timeline.zh.md @@ -13,7 +13,7 @@ ## 编写带追踪的 kernel { #write-a-traced-kernel } -下面是一个完整的、接入了追踪的 warp specialization GEMM（示意用，单缓冲）。编号 `(1)`–`(7)` 的标记是仅有的与追踪相关的代码，各自的说明和 API 文档链接见代码下方。生产环境中的多 stage 版本见 [`src/tileops/kernels/gemm/dense.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/gemm/dense.py)。 +下面是一个完整的、接入了追踪的 warp specialization GEMM（示意用，单缓冲）。编号 `(1)`–`(7)` 的标记是仅有的与追踪相关的代码，各自的说明和 API 文档链接见代码下方。生产环境中的多 stage 版本见 [`src/tileops/_kernels/gemm/dense.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/_kernels/gemm/dense.py)。 ```{ .python .annotate } import functools @@ -62,12 +62,12 @@ def build_gemm(m, n, k, dtype="float16", traced=False): ``` 1. 导入 trace 命名空间。下面的每个调用都是这个 `trace` 对象上的方法，完整说明见 [API 参考](../api/trace.md)。 -1. [`trace.out_idx(n_outputs, traced)`](../api/trace.md#tileops.trace.api._Trace.out_idx) 给出 `@tilelang.jit` 的 `out_idx`。只有 `traced` 时它才多出一个位置给末尾的 `slots` 输出，因此同一个构建函数在开启和关闭追踪时都能使用。 -1. [`trace.group(name, lead)`](../api/trace.md#tileops.trace.api._Trace.group) 声明由哪个 warpgroup 记录。`lead` 是被选出的写入线程（`tx == lead`）：计算仍在所有线程上执行，只有时间戳由 `lead` 写入。 -1. [`trace.range(name, lane)`](../api/trace.md#tileops.trace.api._Trace.range) 是一个 `with` 块，从进入计时到退出，在子行 `lane` 上画成一段条形。控制流放不进 `with` 时，用 [`trace.range_start`](../api/trace.md#tileops.trace.api._Trace.range_start) 与 [`trace.range_end`](../api/trace.md#tileops.trace.api._Trace.range_end)；只需要一个零宽度的时间点时，用 [`trace.record`](../api/trace.md#tileops.trace.api._Trace.record)。 +1. [`trace.out_idx(n_outputs, traced)`](../api/trace.md#tileops.trace._api._Trace.out_idx) 给出 `@tilelang.jit` 的 `out_idx`。只有 `traced` 时它才多出一个位置给末尾的 `slots` 输出，因此同一个构建函数在开启和关闭追踪时都能使用。 +1. [`trace.group(name, lead)`](../api/trace.md#tileops.trace._api._Trace.group) 声明由哪个 warpgroup 记录。`lead` 是被选出的写入线程（`tx == lead`）：计算仍在所有线程上执行，只有时间戳由 `lead` 写入。 +1. [`trace.range(name, lane)`](../api/trace.md#tileops.trace._api._Trace.range) 是一个 `with` 块，从进入计时到退出，在子行 `lane` 上画成一段条形。控制流放不进 `with` 时，用 [`trace.range_start`](../api/trace.md#tileops.trace._api._Trace.range_start) 与 [`trace.range_end`](../api/trace.md#tileops.trace._api._Trace.range_end)；只需要一个零宽度的时间点时，用 [`trace.record`](../api/trace.md#tileops.trace._api._Trace.record)。 1. lane 的名字（`"tma"`、`"barrier"`、`"wgmma"`，以及默认的 `"main"`）就是时间线上的各行。 -1. [`trace.dag(src, dst)`](../api/trace.md#tileops.trace.api._Trace.dag) 声明从一个命名区间指向另一个区间的依赖箭头（`arrive` → `wait`），每出现一次画一条。 -1. [`trace.finalize(func, traced, max_events)`](../api/trace.md#tileops.trace.api._Trace.finalize) 在 `traced` 时 lower 标记并加上 `slots` 输出，否则把标记 strip 成零开销。`traced` **必须是构建函数缓存键的一部分**，这样同一形状下带追踪和不带追踪的两次构建才不会冲突。 +1. [`trace.dag(src, dst)`](../api/trace.md#tileops.trace._api._Trace.dag) 声明从一个命名区间指向另一个区间的依赖箭头（`arrive` → `wait`），每出现一次画一条。 +1. [`trace.finalize(func, traced, max_events)`](../api/trace.md#tileops.trace._api._Trace.finalize) 在 `traced` 时 lower 标记并加上 `slots` 输出，否则把标记 strip 成零开销。`traced` **必须是构建函数缓存键的一部分**，这样同一形状下带追踪和不带追踪的两次构建才不会冲突。 ## 运行带追踪的 kernel { #running-a-traced-kernel } @@ -80,8 +80,8 @@ def forward(self, a, b): return trace.run(compiled, (a, b), stem="gemm_128x256x512") # (2)! ``` -1. 按开关构建对应的 kernel：[`trace.enabled`](../api/trace.md#tileops.trace.api._Trace.enabled) 决定选用带追踪的版本还是 strip 后的版本，它也是标记 `(7)` 所说的缓存键的一部分。 -1. [`trace.run(compiled, inputs, stem=...)`](../api/trace.md#tileops.trace.api._Trace.run) 运行 kernel。带追踪构建的 kernel，末尾的 `slots` 总会被拆出来；开启追踪时，它被解码并写进 `debug/.html`，每次调用都生成一个不会重名的新文件。它内部由 [`decode`](../api/trace.md#tileops.trace.api._Trace.decode) 和 [`dump`](../api/trace.md#tileops.trace.api._Trace.dump) 组成，两者也可以直接调用。 +1. 按开关构建对应的 kernel：[`trace.enabled`](../api/trace.md#tileops.trace._api._Trace.enabled) 决定选用带追踪的版本还是 strip 后的版本，它也是标记 `(7)` 所说的缓存键的一部分。 +1. [`trace.run(compiled, inputs, stem=...)`](../api/trace.md#tileops.trace._api._Trace.run) 运行 kernel。带追踪构建的 kernel，末尾的 `slots` 总会被拆出来；开启追踪时，它被解码并写进 `debug/.html`，每次调用都生成一个不会重名的新文件。它内部由 [`decode`](../api/trace.md#tileops.trace._api._Trace.decode) 和 [`dump`](../api/trace.md#tileops.trace._api._Trace.dump) 组成，两者也可以直接调用。 ## 开启追踪 { #enabling-tracing } @@ -94,7 +94,7 @@ trace.enable() # (1)! c = op(a, b) # (2)! ``` -1. [`trace.enable(output="debug")`](../api/trace.md#tileops.trace.api._Trace.enable) 打开追踪，并指定输出目录，默认是 `debug/`（已被 gitignore）。这个开关只在当前进程内有效：不读取环境变量，也不对 `tilelang` 做 monkeypatch。相关的还有 [`trace.disable()`](../api/trace.md#tileops.trace.api._Trace.disable)、[`trace.enabled`](../api/trace.md#tileops.trace.api._Trace.enabled) 与 [`trace.output`](../api/trace.md#tileops.trace.api._Trace.output)。 +1. [`trace.enable(output="debug")`](../api/trace.md#tileops.trace._api._Trace.enable) 打开追踪，并指定输出目录，默认是 `debug/`（已被 gitignore）。这个开关只在当前进程内有效：不读取环境变量，也不对 `tilelang` 做 monkeypatch。相关的还有 [`trace.disable()`](../api/trace.md#tileops.trace._api._Trace.disable)、[`trace.enabled`](../api/trace.md#tileops.trace._api._Trace.enabled) 与 [`trace.output`](../api/trace.md#tileops.trace._api._Trace.output)。 1. 此后运行的每个带追踪的 kernel 都会写出 `debug/.html`。 在 pytest 中，`--trace-kernel` 会在任何 kernel 构建之前，由 `pytest_configure` 调用 `trace.enable()`： diff --git a/docs/torch-compile.md b/docs/torch-compile.md index cbdf5828..5c615aeb 100644 --- a/docs/torch-compile.md +++ b/docs/torch-compile.md @@ -119,7 +119,7 @@ tracing, graph breaks and guards it refers to are described in `RMSNormFwdOp`'s skeleton, with 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): +[`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): diff --git a/docs/torch-compile.zh.md b/docs/torch-compile.zh.md index 00a8f6d2..eee9b0a9 100644 --- a/docs/torch-compile.zh.md +++ b/docs/torch-compile.zh.md @@ -77,7 +77,7 @@ block(x, w) 本节说明有编译边界的 op 需要编写的代码、边界与 fake 如何生成，以及 target 判定为什么要在节点内部执行。其中涉及的追踪、切图与 guard 见 [dynamo 是怎么工作的](#dynamo)。 -下面是 `RMSNormFwdOp` 的骨架，略去 docstring，完整代码见 [`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): diff --git a/docs/user-guide/benchmark/writing.md b/docs/user-guide/benchmark/writing.md index 2da3ad15..19ad5623 100644 --- a/docs/user-guide/benchmark/writing.md +++ b/docs/user-guide/benchmark/writing.md @@ -34,12 +34,12 @@ This example compares TileOps batched matmul with vLLM, SGLang, and PyTorch/cuBL import pytest from benchmarks import api as bench -from tileops.ops import BmmFwdOp +from tileops.gemm import BMMFwdOp -@pytest.mark.parametrize("case", bench.cases(BmmFwdOp), ids=lambda case: case.id) +@pytest.mark.parametrize("case", bench.cases(BMMFwdOp), ids=lambda case: case.id) def test_bmm_bench(case, vllm_bmm, sglang_bmm): - op = BmmFwdOp(**case.arguments) + op = BMMFwdOp(**case.arguments) bench.Runner(op, case).compare( { @@ -52,7 +52,7 @@ def test_bmm_bench(case, vllm_bmm, sglang_bmm): Here is how the example works: -- `bench.cases(BmmFwdOp)` creates a case for each dtype case of every manifest workload +- `bench.cases(BMMFwdOp)` creates a case for each dtype case of every manifest workload row. The `case.id` also serves as the pytest test ID. - `case.arguments` supplies the arguments used to construct the op. - `vllm_bmm` and `sglang_bmm` are callables that accept the same inputs as the op. @@ -82,8 +82,8 @@ just one line. Here is an excerpt from `_cases/gemm.py`: ```python ENTRIES = { - "BmmFwdOp": Entry(BmmWorkload.from_call), - "GemmFwdOp": Entry(GemmWorkload.from_call), + "BMMFwdOp": Entry(BMMWorkload.from_call), + "GEMMFwdOp": Entry(GEMMWorkload.from_call), } ``` diff --git a/docs/user-guide/benchmark/writing.zh.md b/docs/user-guide/benchmark/writing.zh.md index 110952b8..4c1114c8 100644 --- a/docs/user-guide/benchmark/writing.zh.md +++ b/docs/user-guide/benchmark/writing.zh.md @@ -31,12 +31,12 @@ bench.Runner(op, case).compare(implementations) import pytest from benchmarks import api as bench -from tileops.ops import BmmFwdOp +from tileops.gemm import BMMFwdOp -@pytest.mark.parametrize("case", bench.cases(BmmFwdOp), ids=lambda case: case.id) +@pytest.mark.parametrize("case", bench.cases(BMMFwdOp), ids=lambda case: case.id) def test_bmm_bench(case, vllm_bmm, sglang_bmm): - op = BmmFwdOp(**case.arguments) + op = BMMFwdOp(**case.arguments) bench.Runner(op, case).compare( { @@ -49,8 +49,8 @@ def test_bmm_bench(case, vllm_bmm, sglang_bmm): 这段代码中： -- `bench.cases(BmmFwdOp)` 为 manifest 中每个 workload 行的每个 dtype case 各生成一个 case，`case.id` 同时用作 pytest 的用例 id。 -- `BmmFwdOp(**case.arguments)` 用 case 提供的构造参数创建 op。 +- `bench.cases(BMMFwdOp)` 为 manifest 中每个 workload 行的每个 dtype case 各生成一个 case，`case.id` 同时用作 pytest 的用例 id。 +- `BMMFwdOp(**case.arguments)` 用 case 提供的构造参数创建 op。 - `vllm_bmm` 与 `sglang_bmm` 是已经适配成与 op 相同输入签名的 callable，`Runner` 会以 `case.inputs` 调用它们。 - 映射中的键会原样写入报告。`tileops` 这个名字已经由 op 占用，不能再用作键；如果只传入一个实现而不是映射，它在报告中记为 `baseline`。 - `case.reference` 以 `"torch-cublas"` 为名传入，因此它既负责生成期望结果，也作为一个对比实现参与计时。 @@ -70,8 +70,8 @@ python -m pytest benchmarks/ops/bench_bmm.py --tileops-verify ```python ENTRIES = { - "BmmFwdOp": Entry(BmmWorkload.from_call), - "GemmFwdOp": Entry(GemmWorkload.from_call), + "BMMFwdOp": Entry(BMMWorkload.from_call), + "GEMMFwdOp": Entry(GEMMWorkload.from_call), } ``` diff --git a/docs/user-guide/development.md b/docs/user-guide/development.md index a96e8731..b83962b3 100644 --- a/docs/user-guide/development.md +++ b/docs/user-guide/development.md @@ -84,7 +84,7 @@ updates the spec. ### Kernel compilation and caching { #compile } ```python -op = GemmFwdOp() +op = GEMMFwdOp() d = op(a, b) # first call: compiles the kernel and caches it d = op(a, b) # the same call: reuses the cache ``` @@ -148,7 +148,7 @@ python -m pytest -q | `[Type][foundry][Scope] ` | `[Perf][foundry][Elementwise] Take the floored tier's reciprocal from one MUFU instruction` | - CI checks the PR title's format. The values of `Type` are defined in - [`.claude/conventions/types.sh`](https://github.com/tile-ai/TileOPs/blob/main/.claude/conventions/types.sh). + [`scripts/ci/types.sh`](https://github.com/tile-ai/TileOPs/blob/main/scripts/ci/types.sh). - `foundry` marks a PR whose kernels were generated by [TileFoundry](https://github.com/tile-ai/TileFoundry). diff --git a/docs/user-guide/development.zh.md b/docs/user-guide/development.zh.md index 895303ee..1fea002a 100644 --- a/docs/user-guide/development.zh.md +++ b/docs/user-guide/development.zh.md @@ -71,7 +71,7 @@ python -m pytest -q tests -m smoke ### kernel 的编译与缓存 { #compile } ```python -op = GemmFwdOp() +op = GEMMFwdOp() d = op(a, b) # 首次调用：编译 kernel 并写入缓存 d = op(a, b) # 相同的调用：复用缓存 ``` @@ -121,7 +121,7 @@ python -m pytest -q | `[Type][Scope] ` | `[BugFix][Elementwise] Build the floored tiers at every tuned fold width` | | `[Type][foundry][Scope] ` | `[Perf][foundry][Elementwise] Take the floored tier's reciprocal from one MUFU instruction` | -- CI 检查 PR 标题的格式。`Type` 的取值定义在 [`.claude/conventions/types.sh`](https://github.com/tile-ai/TileOPs/blob/main/.claude/conventions/types.sh) 中。 +- CI 检查 PR 标题的格式。`Type` 的取值定义在 [`scripts/ci/types.sh`](https://github.com/tile-ai/TileOPs/blob/main/scripts/ci/types.sh) 中。 - `foundry` 标记 kernel 由 [TileFoundry](https://github.com/tile-ai/TileFoundry) 生成的 PR。 ### 描述 { #pr-body } diff --git a/docs/user-guide/dispatch/backends.md b/docs/user-guide/dispatch/backends.md index 72e0e2be..58e9182d 100644 --- a/docs/user-guide/dispatch/backends.md +++ b/docs/user-guide/dispatch/backends.md @@ -1,78 +1,24 @@ # How a backend joins TileOPs -A third-party backend has two ways to join, ordered from the smaller range of calls taken over to the larger: add an implementation to a kernel interface, or replace a whole op. In the first, the backend's class and the in-tree implementations follow the same kernel interface; how to write one is in [Adding a kernel to an op](writing.md). +A third-party backend joins through a target: it registers a builder for an op and takes over every call of that op. The kernel interfaces and call specs of TileOPs' ops are not public API, so a backend does not add an implementation to a single kernel interface. -## 1. Choose the way by the range of calls to take over {#choose} +## 1. What a target serves and what the in-tree implementations serve {#choose} -![The ranges taken over by the two ways of extending](img/extension.svg) +![The ranges served by a target and by the in-tree implementations](img/extension.svg) -**Figure 1** The range each of the two ways of extending takes over. Teal marks the parts provided by the system, purple the parts written by TileOPs developers, and green the parts provided by the backend. +**Figure 1** The range a target serves and the range the in-tree implementations serve. Teal marks the parts provided by the system, purple the parts written by TileOPs developers, and green the parts provided by the backend. -**Table 1** Comparison of the two ways +**Table 1** What a target covers -| No. | Item | `register_kernel_type` | target | -| --- | --- | --- | --- | -| 1 | What changes | a new key, with its own applicability and precedence | every call of the op, except calls whose written tensors are all empty | -| 2 | Applies to | every instance of the op constructed after registration | the op instances that select the target | -| 3 | Contract followed | the kernel interface | the op's signature in the manifest | -| 4 | Calls the new class does not serve | still served by the in-tree implementations | none; the target serves every call, except calls whose written tensors are all empty | +| No. | Item | target | +| --- | --- | --- | +| 1 | What changes | every call of the op, except calls whose written tensors are all empty | +| 2 | Applies to | the op instances that select the target | +| 3 | Contract followed | the op's signature in the manifest | -## 2. Add an implementation: register_kernel_type {#register} +## 2. Replace a whole op: target {#target} -`tileops.backend.register_kernel_type(op, key, kernel_type)` adds an implementation to an op. It is the backend form of step 1 in [Adding a kernel to an op](writing.md). `op` is the op's class name, `key` is the new implementation's name, and `kernel_type` is a class that subclasses both `Kernel` and one of the op's kernel interfaces; the inheritance decides which interface it belongs to, and the class is built through its own class method `entry_for(call)`. Step 2 is the same as for in-tree implementations. - -```python -# tests/test_kernel_dispatch.py -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 - ) - - -class _NarrowTorchLayerNorm(_TorchLayerNorm): - """An added implementation for short rows, which wins over the in-tree one there.""" - - preferred_over = frozenset({"layer_norm"}) - - @classmethod - def refusal(cls, call: LayerNormCall) -> "str | None": - return None if call.n <= 256 else "serves rows of at most 256" - -register_kernel_type("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerNorm) -``` - -- The new implementation overlaps `LayerNormKernel` on `n <= 256`, and `LayerNormKernel` is not general, so the new implementation declares `preferred_over`. -- Calls the new implementation does not serve, such as `n = 1024`, are still served by the in-tree implementation. -- The new implementation enters only op instances constructed after registration. -- The new implementation declares the device types it runs on in `devices`. On a device type other than CUDA, the backend registers the device facts with `register_device_facts`; see [Adding a hardware backend § Device facts](../../backends.md#device-facts). -- Registering the same key twice under one op raises `BackendError`. A key equal to an in-tree key raises `reuse keys it has` when the op is constructed. -- When the new implementation inherits none of the op's kernel interfaces, states `entry_for` other than as a class method, or has a `forward` that does not take the interface's arguments, constructing the op raises; see [Adding a kernel to an op § 5](writing.md#tests). - -Registration happens when the backend module is imported. The backend declares an entry point in `pyproject.toml`, and TileOPs imports it when the first op is constructed: - -```toml -[project.entry-points."tileops.backends"] -acme = "tileops_acme" -``` - -When the module fails to import, all its registrations are undone and TileOPs emits a `RuntimeWarning`. The failure records are available through `tileops.backend.load_failures()`. - -## 3. Replace a whole op: target {#target} - -The full target protocol, a runnable template backend, and common errors are in [Adding a hardware backend](../../backends.md). This section only summarizes how a target differs from the other way. +The full target protocol, a runnable template backend, and common errors are in [Adding a hardware backend](../../backends.md). This section only summarizes it. A target is a name a backend gives to a set of kernels. Once a target registers a builder for an op, every call of an op instance that selects the target is served by the target, and the op's own `forward` does not run. The exception is a call whose written tensors are all empty: then neither the target nor the in-tree implementation runs, and the outputs are constructed from the signature. @@ -94,3 +40,12 @@ register_kernel_builder(op="RMSNormFwdOp", target="acme", build_kernel=build_rms - An op instance fixes its target on the first call that needs to run an implementation, and the target does not change afterwards. The construction argument `target=` comes first, then the process default set by `set_default_target`, and finally each target's detector inspects the call's device. When no target claims the device, the in-tree implementation is used; when several targets claim it, the call raises `AmbiguousTargetError`. A call whose written tensors are all empty does not fix the target. When the call that fixes the target fails, the choice is undone and the next call fixes it again. - A call passed to a target has already passed the checks generated from the signature. Every tensor is on the call's device, except tensors declared `device: cpu`. Every input that is not written is contiguous, and a written input is contiguous when it declares `contiguous: true`. - When the selected target has registered no builder for this op and the op holds kernels of its own, the call raises `OpNotAvailableError` and does not fall back to the in-tree implementation. + +Registration happens when the backend module is imported. The backend declares an entry point in `pyproject.toml`, and TileOPs imports it when the first op is constructed: + +```toml +[project.entry-points."tileops.backends"] +acme = "tileops_acme" +``` + +When the module fails to import, all its registrations are undone and TileOPs emits a `RuntimeWarning`. The failure records are available through `tileops.backend.load_failures()`. diff --git a/docs/user-guide/dispatch/backends.zh.md b/docs/user-guide/dispatch/backends.zh.md index 44614c9f..8b8c2c06 100644 --- a/docs/user-guide/dispatch/backends.zh.md +++ b/docs/user-guide/dispatch/backends.zh.md @@ -1,78 +1,24 @@ # backend 如何接入 TileOPs -第三方 backend 按要接管的调用范围，从小到大有两种方式：为一个 kernel 接口新增实现、替换整个 op。第一种方式中 backend 的类与 in-tree 实现遵守同一份 kernel 接口，写法见[如何为 op 新增 kernel](writing.md)。 +第三方 backend 通过 target 接入：为 op 注册 builder，接管这个 op 的全部调用。TileOPs 各 op 的 kernel 接口与 call spec 不是公开 API，backend 不为单个 kernel 接口新增实现。 -## 1. 按要接管的调用范围选择接入方式 {#choose} +## 1. target 与 in-tree 实现各自服务的范围 {#choose} -![两种扩展方式接管的范围](img/extension.svg) +![target 与 in-tree 实现服务的范围](img/extension.svg) -**图 1** 两种扩展方式各自接管的范围。青色为系统提供的部分，紫色为 TileOPs 开发者编写的部分，绿色为 backend 提供的部分。 +**图 1** target 与 in-tree 实现各自服务的范围。青色为系统提供的部分，紫色为 TileOPs 开发者编写的部分，绿色为 backend 提供的部分。 -**表 1** 两种方式的比较 +**表 1** target 的作用范围 -| No. | 项目 | `register_kernel_type` | target | -| --- | --- | --- | --- | -| 1 | 改变什么 | 新增一个 key，带有它自己的适用范围与优先关系 | op 的全部调用，所有写入张量都为空的调用除外 | -| 2 | 作用于 | 注册之后构造的所有该 op 实例 | 选中该 target 的 op 实例 | -| 3 | 依据的契约 | kernel 接口 | manifest 中 op 的签名 | -| 4 | 新类不服务的调用 | 仍由 in-tree 实现服务 | 不存在，target 服务全部调用，所有写入张量都为空的调用除外 | +| No. | 项目 | target | +| --- | --- | --- | +| 1 | 改变什么 | op 的全部调用，所有写入张量都为空的调用除外 | +| 2 | 作用于 | 选中该 target 的 op 实例 | +| 3 | 依据的契约 | manifest 中 op 的签名 | -## 2. 新增实现：register_kernel_type {#register} +## 2. 替换整个 op：target {#target} -`tileops.backend.register_kernel_type(op, key, kernel_type)` 为 op 新增一个实现，是[如何为 op 新增 kernel](writing.md) 第一步的 backend 形式。`op` 是 op 的类名，`key` 是新实现的名字，`kernel_type` 是同时继承 `Kernel` 与 op 的某个 kernel 接口的类，属于哪个接口由继承关系决定，由它自己的类方法 `entry_for(call)` 构建。第二步与 in-tree 实现相同。 - -```python -# tests/test_kernel_dispatch.py -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 - ) - - -class _NarrowTorchLayerNorm(_TorchLayerNorm): - """An added implementation for short rows, which wins over the in-tree one there.""" - - preferred_over = frozenset({"layer_norm"}) - - @classmethod - def refusal(cls, call: LayerNormCall) -> "str | None": - return None if call.n <= 256 else "serves rows of at most 256" - -register_kernel_type("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerNorm) -``` - -- 新实现在 `n <= 256` 上与 `LayerNormKernel` 重叠，后者不是 general，因此新实现声明 `preferred_over`。 -- 新实现不服务的调用仍由 in-tree 实现服务，例如 `n = 1024`。 -- 新实现只进入注册之后构造的 op 实例。 -- 新实现在 `devices` 中声明它运行的设备类型。CUDA 以外的设备类型，其设备事实由 backend 用 `register_device_facts` 注册，见[接入新的硬件 backend § 设备事实](../../backends.md#device-facts)。 -- 同一个 op 下重复注册同一个 key 报 `BackendError`；key 与 in-tree 的 key 相同时，构造 op 时报 `reuse keys it has`。 -- 新实现不继承 op 的任何 kernel 接口、`entry_for` 不是类方法、或 `forward` 不接受接口的参数时，构造 op 时报错，见[如何为 op 新增 kernel § 5](writing.md#tests)。 - -注册在 backend 模块被导入时完成。backend 在 `pyproject.toml` 中声明 entry point，TileOPs 在构造第一个 op 时导入它： - -```toml -[project.entry-points."tileops.backends"] -acme = "tileops_acme" -``` - -模块导入失败时，它的全部注册被撤销，TileOPs 发出 `RuntimeWarning`，失败记录可以通过 `tileops.backend.load_failures()` 查看。 - -## 3. 替换整个 op：target {#target} - -target 协议的完整写法、一个可运行的模板 backend 与常见报错，见[接入新的硬件 backend](../../backends.md)。本节只概述它与另一种方式的区别。 +target 协议的完整写法、一个可运行的模板 backend 与常见报错，见[接入新的硬件 backend](../../backends.md)。本节只给出概要。 target 是 backend 为一组 kernel 取的名字。target 为某个 op 注册 builder 之后，选中该 target 的 op 实例的调用都由 target 服务，op 自己的 `forward` 不执行。一次调用写入的所有张量都为空时例外：这时 target 与 in-tree 实现都不执行，输出按签名构造。 @@ -94,3 +40,12 @@ register_kernel_builder(op="RMSNormFwdOp", target="acme", build_kernel=build_rms - op 实例在第一次需要执行实现的调用中确定 target，此后不变：构造参数 `target=` 优先，其次是 `set_default_target` 设置的进程默认值，最后由各 target 的 detector 检测调用设备；没有 target 认领时使用 in-tree 实现，多个 target 认领时报 `AmbiguousTargetError`。写入张量都为空的调用不确定 target；确定 target 的那次调用失败时，选择被撤销，下一次调用重新确定。 - 传给 target 的调用已经通过由签名生成的检查；除声明了 `device: cpu` 的张量外，所有张量位于调用设备上；所有不被写入的输入都是连续的，被写入的输入在声明了 `contiguous: true` 时是连续的。 - 选中的 target 没有为这个 op 注册 builder、而 op 自己持有 kernel 时，调用报 `OpNotAvailableError`，不退回 in-tree 实现。 + +注册在 backend 模块被导入时完成。backend 在 `pyproject.toml` 中声明 entry point，TileOPs 在构造第一个 op 时导入它： + +```toml +[project.entry-points."tileops.backends"] +acme = "tileops_acme" +``` + +模块导入失败时，它的全部注册被撤销，TileOPs 发出 `RuntimeWarning`，失败记录可以通过 `tileops.backend.load_failures()` 查看。 diff --git a/docs/user-guide/dispatch/img/extension.svg b/docs/user-guide/dispatch/img/extension.svg index 0bfae41e..655978ee 100644 --- a/docs/user-guide/dispatch/img/extension.svg +++ b/docs/user-guide/dispatch/img/extension.svg @@ -1 +1 @@ -target: replaces the whole opThe op instance's in-tree implementationImplementations of one kernel interface(installed at construction)build_kernel registered byregister_kernel_builderThe Op class's forwardClasses in kernel_typesinheriting the interfaceregister_kernel_typeadds a key(every instance built afterwards)Generated call checksa target registereda builder for the opserved in-treekernel_for selects amongthese by the selection ruleCyan: provided by the systemPurple: written by TileOPs developersGreen: provided by a backend +target: replaces the whole opThe op instance's in-tree implementationImplementations of one kernel interface(installed at construction)build_kernel registered byregister_kernel_builderThe Op class's forwardClasses in kernel_typesinheriting the interfaceGenerated call checksa target registereda builder for the opserved in-treekernel_for selects amongthese by the selection ruleCyan: provided by the systemPurple: written by TileOPs developersGreen: provided by a backend diff --git a/docs/user-guide/dispatch/img/extension.zh.svg b/docs/user-guide/dispatch/img/extension.zh.svg index a4d09ad0..babbb891 100644 --- a/docs/user-guide/dispatch/img/extension.zh.svg +++ b/docs/user-guide/dispatch/img/extension.zh.svg @@ -1 +1 @@ -target：替换整个 opop 实例的 in-tree 实现一个 kernel 接口的实现（构造时安装）register_kernel_builder注册的 build_kernelOp 类的 forwardkernel_types 中继承该接口的类register_kernel_type新增的 key（此后构造的所有实例）生成的调用检查target 为该 op注册了 builder由 in-tree 实现服务kernel_for 按选择规则在这些实现中选择青色：系统提供紫色：TileOPs 开发者编写绿色：backend 提供 +target：替换整个 opop 实例的 in-tree 实现一个 kernel 接口的实现（构造时安装）register_kernel_builder注册的 build_kernelOp 类的 forwardkernel_types 中继承该接口的类生成的调用检查target 为该 op注册了 builder由 in-tree 实现服务kernel_for 按选择规则在这些实现中选择青色：系统提供紫色：TileOPs 开发者编写绿色：backend 提供 diff --git a/docs/user-guide/dispatch/index.md b/docs/user-guide/dispatch/index.md index 9c549c35..0388df4a 100644 --- a/docs/user-guide/dispatch/index.md +++ b/docs/user-guide/dispatch/index.md @@ -1,11 +1,11 @@ # How an op selects a kernel -A TileOPs op can have several kernels, and each kernel serves part of the op's calls. This guide describes how an op selects a kernel on a call, and what a developer writes to add a kernel. It is written for two readers: developers who add a kernel or an op to TileOPs, and authors of third-party backends who add kernels or take over a whole op. The reasons behind each design choice are in the design document [Op Interface Design § Kernel selection](../../design/ops-design.md#kernel-selection). +A TileOPs op can have several kernels, and each kernel serves part of the op's calls. This guide describes how an op selects a kernel on a call, and what a developer writes to add a kernel. It is written for two readers: developers who add a kernel or an op to TileOPs, and authors of third-party backends who take over a whole op. The reasons behind each design choice are in the design document [Op Interface Design § Kernel selection](../../design/ops-design.md#kernel-selection). This page describes the selection process and its terms. The other two pages are: - [Adding a kernel to an op](writing.md): the two steps, the default rules, opening a new kernel interface, tests, and common errors; -- [How a backend joins TileOPs](backends.md): `register_kernel_type` and targets. +- [How a backend joins TileOPs](backends.md): replacing a whole op with a target. ## 1. Who selects the kernel: the op selects an interface, dispatch selects an implementation {#two-levels} @@ -16,7 +16,7 @@ The op first selects a **kernel interface** by semantics. Dispatch then selects | No. | Level | Decided by | Based on | Written in | | --- | --- | --- | --- | --- | | 1 | Select a kernel interface | the op | semantics, that is, where the call contracts differ, such as training versus inference | the op's `forward`, which names the interface when it calls `kernel_for` | -| 2 | Select an implementation | dispatch, shared by all ops | the availability, applicability, and precedence each implementation declares | the selection algorithm is in `src/tileops/ops/op_base.py`; each implementation's declarations are on its own kernel class | +| 2 | Select an implementation | dispatch, shared by all ops | the availability, applicability, and precedence each implementation declares | the selection algorithm is in `src/tileops/_ops/op_base.py`; each implementation's declarations are on its own kernel class | Take `BatchNormFwdOp` as an example. The op selects the interface `batch_norm_fwd_train` or `batch_norm_fwd_infer` by its construction argument `training`. The training interface has four implementations, and dispatch selects one of them by batch size, channel count, and spatial size. @@ -25,9 +25,9 @@ Take `BatchNormFwdOp` as an example. The op selects the interface `batch_norm_fw On every call, the op builds a **call spec**, which records the facts of this call that affect selection and building. It then calls `self.kernel_for(interface, call)`: ```python -# src/tileops/ops/norm/layer_norm.py -call = LayerNormCall(device=x.device, n=math.prod(ns), eps=self.eps, dtype=x.dtype) -return self.kernel_for("layer_norm", call)(x, weight, bias) +# src/tileops/_ops/norm/layer_norm.py +call = LayerNormCall(device=input.device, n=math.prod(ns), eps=self.eps, dtype=input.dtype) +return self.kernel_for("layer_norm", call)(input, weight, bias) ``` ![The path of one kernel_for call](img/call.svg) @@ -45,18 +45,18 @@ When no key can run on the call's device type, the call raises `in-tree kernels ## 3. Two things to do when adding a kernel {#hooks} -To add a kernel to an existing kernel interface, TileOPs developers and backend authors both do the following two things: +To add a kernel to an existing kernel interface, a TileOPs developer does the following two things: **Table 2** The two steps of adding a kernel -| No. | Step | TileOPs developer | Backend author | Described in | -| --- | --- | --- | --- | --- | -| 1 | Register the implementation | the class subclasses the kernel interface and is added to the op's `kernel_types` | the class subclasses the kernel interface; call `register_kernel_type(op, key, cls)` | [Adding a kernel to an op § 1](writing.md#register) | -| 2 | Declare which calls the implementation serves | write `refusal`; write `preferred_over` only when it overlaps another non-general implementation | same as the TileOPs developer | [Adding a kernel to an op § 2](writing.md#rule) | +| No. | Step | What to write | Described in | +| --- | --- | --- | --- | +| 1 | Register the implementation | the class subclasses the kernel interface and is added to the op's `kernel_types` | [Adding a kernel to an op § 1](writing.md#register) | +| 2 | Declare which calls the implementation serves | write `refusal`; write `preferred_over` only when it overlaps another non-general implementation | [Adding a kernel to an op § 2](writing.md#rule) | An implementation that declares none of `devices`, `supported_archs`, `refusal`, `general`, and `preferred_over` is available on CUDA devices of every architecture, serves every call, and has no precedence relation with other implementations. An interface with a single implementation therefore only needs the class to subclass the interface. -A backend can also replace a whole op; see [How a backend joins TileOPs](backends.md). +A backend does not add a single kernel; it replaces a whole op, see [How a backend joins TileOPs](backends.md). ## 4. Terms used in this guide {#terms} diff --git a/docs/user-guide/dispatch/index.zh.md b/docs/user-guide/dispatch/index.zh.md index 39a45926..a1557c76 100644 --- a/docs/user-guide/dispatch/index.zh.md +++ b/docs/user-guide/dispatch/index.zh.md @@ -1,11 +1,11 @@ # op 如何选择 kernel -一个 TileOPs op 可以有多个 kernel，每个 kernel 服务一部分调用。本指南说明 op 在一次调用中如何选中 kernel，以及新增一个 kernel 要写什么。本指南面向两类读者：在 TileOPs 中新增 kernel 或 op 的开发者，以及为第三方 backend 新增 kernel 或接管整个 op 的作者。各项设计的理由见设计文档 [Op Interface Design § Kernel selection](../../design/ops-design.md#kernel-selection)。 +一个 TileOPs op 可以有多个 kernel，每个 kernel 服务一部分调用。本指南说明 op 在一次调用中如何选中 kernel，以及新增一个 kernel 要写什么。本指南面向两类读者：在 TileOPs 中新增 kernel 或 op 的开发者，以及用第三方 backend 接管整个 op 的作者。各项设计的理由见设计文档 [Op Interface Design § Kernel selection](../../design/ops-design.md#kernel-selection)。 本页说明选择的过程与术语，另外两页分别是： - [如何为 op 新增 kernel](writing.md)：两步的写法、默认规则、新开 kernel 接口、测试与常见报错； -- [backend 如何接入 TileOPs](backends.md)：`register_kernel_type` 与 target。 +- [backend 如何接入 TileOPs](backends.md)：用 target 替换整个 op。 ## 1. 谁来选 kernel：op 选接口，派发机制选实现 {#two-levels} @@ -16,7 +16,7 @@ op 先按语义选定一个 **kernel 接口**，派发机制再按调用的 shap | No. | 层 | 由谁决定 | 依据 | 写在哪里 | | --- | --- | --- | --- | --- | | 1 | 选 kernel 接口 | op | 语义，即调用契约不同之处，例如训练与推理 | op 的 `forward`，调用 `kernel_for` 时给出接口名 | -| 2 | 选实现 | 派发机制，所有 op 共用 | 每个实现声明的可用性、适用范围与优先关系 | 选择算法在 `src/tileops/ops/op_base.py`；各实现的声明写在自己的 kernel 类上 | +| 2 | 选实现 | 派发机制，所有 op 共用 | 每个实现声明的可用性、适用范围与优先关系 | 选择算法在 `src/tileops/_ops/op_base.py`；各实现的声明写在自己的 kernel 类上 | 以 `BatchNormFwdOp` 为例：op 按构造参数 `training` 选择接口 `batch_norm_fwd_train` 或 `batch_norm_fwd_infer`；训练接口有四个实现，派发机制按 batch、channel 数与空间尺寸选中其中一个。 @@ -25,9 +25,9 @@ op 先按语义选定一个 **kernel 接口**，派发机制再按调用的 shap op 每次调用构造一个 **call spec**，它记录这次调用中影响选择与构建的事实，然后调用 `self.kernel_for(interface, call)`： ```python -# src/tileops/ops/norm/layer_norm.py -call = LayerNormCall(device=x.device, n=math.prod(ns), eps=self.eps, dtype=x.dtype) -return self.kernel_for("layer_norm", call)(x, weight, bias) +# src/tileops/_ops/norm/layer_norm.py +call = LayerNormCall(device=input.device, n=math.prod(ns), eps=self.eps, dtype=input.dtype) +return self.kernel_for("layer_norm", call)(input, weight, bias) ``` ![一次 kernel_for 调用的路径](img/call.svg) @@ -45,18 +45,18 @@ return self.kernel_for("layer_norm", call)(x, weight, bias) ## 3. 新增一个 kernel 要做的两件事 {#hooks} -为已有的 kernel 接口新增一个 kernel，TileOPs 开发者与 backend 作者都做以下两件事： +TileOPs 开发者为已有的 kernel 接口新增一个 kernel，做以下两件事： **表 2** 新增一个 kernel 的两步 -| No. | 步骤 | TileOPs 开发者 | backend 作者 | 说明所在 | -| --- | --- | --- | --- | --- | -| 1 | 登记实现 | 类继承 kernel 接口，加入 op 的 `kernel_types` | 类继承 kernel 接口，调用 `register_kernel_type(op, key, cls)` | [新增 kernel 1](writing.md#register) | -| 2 | 声明实现服务哪些调用 | 写 `refusal`；只在与另一个非 general 实现重叠时写 `preferred_over` | 同左 | [新增 kernel 2](writing.md#rule) | +| No. | 步骤 | 做法 | 说明所在 | +| --- | --- | --- | --- | +| 1 | 登记实现 | 类继承 kernel 接口，加入 op 的 `kernel_types` | [新增 kernel 1](writing.md#register) | +| 2 | 声明实现服务哪些调用 | 写 `refusal`；只在与另一个非 general 实现重叠时写 `preferred_over` | [新增 kernel 2](writing.md#rule) | 实现不声明 `devices`、`supported_archs`、`refusal`、`general` 与 `preferred_over` 时，默认在所有架构的 CUDA 设备上可用，服务全部调用，与其他实现没有优先关系。只有一个实现的接口因此只需继承接口。 -backend 还可以替换整个 op，见 [backend 如何接入 TileOPs](backends.md)。 +backend 不新增单个 kernel，而是替换整个 op，见 [backend 如何接入 TileOPs](backends.md)。 ## 4. 本指南使用的术语 {#terms} diff --git a/docs/user-guide/dispatch/writing.md b/docs/user-guide/dispatch/writing.md index 55029da5..49286538 100644 --- a/docs/user-guide/dispatch/writing.md +++ b/docs/user-guide/dispatch/writing.md @@ -4,10 +4,10 @@ This page describes the two steps of adding a kernel to an existing op, the defa ## 1. Step 1: register the implementation with the op {#register} -An implementation is a class that subclasses both `Kernel` (or one of its subclasses) and a kernel interface. A TileOPs developer adds it to the op's `kernel_types` under a snake_case key. A backend author calls `register_kernel_type`; see [How a backend joins TileOPs § 2](backends.md#register). +An implementation is a class that subclasses both `Kernel` (or one of its subclasses) and a kernel interface. A TileOPs developer adds it to the op's `kernel_types` under a snake_case key. ```python -# src/tileops/ops/norm/batch_norm.py +# src/tileops/_ops/norm/batch_norm.py class BatchNormFwdOp(Op): kernel_types: ClassVar[Mapping[str, type[Kernel]]] = { "fwd_train_whole": BatchNormFwdTrainWholeKernel, @@ -51,7 +51,7 @@ When a new implementation's applicability overlaps an existing implementation, w The three GQA dense decode implementations are case 3. bs1 serves calls with batch 1, and long-context serves calls with a long KV. The two overlap when batch is 1 and KV is long, and long-context declares that it wins: ```python -# src/tileops/kernels/attention/gqa/decode.py +# src/tileops/_kernels/attention/gqa/decode.py class GQADecodeLongContextKernel(GQADecodeKernel): general: bool = False preferred_over = frozenset({"gqa_dense_decode_bs1"}) @@ -100,7 +100,7 @@ Two non-general implementations need a precedence declaration as soon as their a An interface with a single implementation therefore only needs the class to subclass the interface. `LayerNormKernel` declares only its supported architectures and its own `entry_for`: ```python -# src/tileops/kernels/norm/layer_norm.py +# src/tileops/_kernels/norm/layer_norm.py class LayerNormKernel(Kernel, LayerNormFwdInterface): supported_archs: list[int] = [80, 86, 89, 90] @@ -123,10 +123,10 @@ A new kernel that changes the semantics or the call contract, for example by ret | 3 | a faster algorithm on some shape range or some architecture | a new implementation | | 4 | the same algorithm with a split count of 1 or greater than 1 | a parameter inside the implementation | -A new interface needs a call spec type and an interface class. When the family has several kernel files, both are written in `src/tileops/kernels//call_spec.py`; when it has a single kernel file, both are written in that kernel file. Interface classes are named `{Name}{Fwd|Bwd}Interface`, with the word for the variant placed before the direction suffix, for example `BatchNormTrainFwdInterface`; the pre-commit hook `interface-names-lint` checks this. Several interfaces of one family can share a call spec type; for example, the six interfaces of BatchNorm and InstanceNorm all use `BatchNormCall`. +A new interface needs a call spec type and an interface class. When the family has several kernel files, both are written in `src/tileops/_kernels//call_spec.py`; when it has a single kernel file, both are written in that kernel file. Interface classes are named `{Name}{Fwd|Bwd}Interface`, with the word for the variant placed before the direction suffix, for example `BatchNormTrainFwdInterface`; the pre-commit hook `interface-names-lint` checks this. Several interfaces of one family can share a call spec type; for example, the six interfaces of BatchNorm and InstanceNorm all use `BatchNormCall`. ```python -# src/tileops/kernels/norm/call_spec.py +# src/tileops/_kernels/norm/call_spec.py @dataclasses.dataclass(frozen=True) class LayerNormCall(CallSpec): """The facts that select an implementation normalizing trailing rows and build it. @@ -161,7 +161,7 @@ class LayerNormFwdInterface(KernelInterface): ``` - The call spec's fields are the call facts that implementations read in `refusal` and `entry_for`: shapes, dtypes, semantic parameters (including those fixed when the op is constructed), and `device`. Fields hold only immutable values. They do not hold tensor contents, tuning policy, or priorities; device facts are provided by `CallSpec`. -- The interface class's `request` points to the call spec type. The parameter list of the abstract method `forward` is the arguments the op passes when it calls the kernel. Its docstring states, for each tensor, the shape, dtype, memory layout, device, and whether it is written in place, and it states the return value. A backend writes its implementation from this contract alone. +- The interface class's `request` points to the call spec type. The parameter list of the abstract method `forward` is the arguments the op passes when it calls the kernel. Its docstring states, for each tensor, the shape, dtype, memory layout, device, and whether it is written in place, and it states the return value. An implementation is written from this contract alone. - Methods other than `forward` that the op calls on the entry are also written as abstract methods of the interface. An implementation must implement every abstract method of the interface; otherwise the class cannot be instantiated. - Values that every implementation must agree on are written as ordinary class methods of the interface and computed by the interface. For example, when the op allocates an output buffer by size before selection, a class method of the interface computes that size from the call spec; the size is not read from the selected implementation. - The docstring states only what the implementation actually does. If it declares the inputs contiguous, the op calls `.contiguous()` before the call. Behavior the implementation does not guarantee (for example, the result when an input contains NaN) is stated as undefined, not as a guarantee. @@ -174,8 +174,8 @@ Write one test case for each applicability range and for each boundary between a ```python # tests/ops/test_family_dispatch.py def test_gemm_k_too_narrow_to_vectorize_is_refused_during_selection() -> None: - op = GemmFwdOp() - call = GemmCall(arch=_SM90, sm_count=132, m=64, n=64, k=1, dtype=torch.float16, trans_b=True) + op = GEMMFwdOp() + call = GEMMCall(arch=_SM90, sm_count=132, m=64, n=64, k=1, dtype=torch.float16, trans_b=True) with pytest.raises(ValueError, match="k must span at least one"): op.key_for("gemm", call) diff --git a/docs/user-guide/dispatch/writing.zh.md b/docs/user-guide/dispatch/writing.zh.md index 375c39ba..e5f6b7d0 100644 --- a/docs/user-guide/dispatch/writing.zh.md +++ b/docs/user-guide/dispatch/writing.zh.md @@ -4,10 +4,10 @@ ## 1. 第一步：把实现登记到 op {#register} -实现是同时继承 `Kernel`（或它的子类）与一个 kernel 接口的类。TileOPs 开发者把它加入 op 的 `kernel_types`，key 使用 snake_case；backend 作者调用 `register_kernel_type`，见 [backend 接入 2](backends.md#register)。 +实现是同时继承 `Kernel`（或它的子类）与一个 kernel 接口的类。TileOPs 开发者把它加入 op 的 `kernel_types`，key 使用 snake_case。 ```python -# src/tileops/ops/norm/batch_norm.py +# src/tileops/_ops/norm/batch_norm.py class BatchNormFwdOp(Op): kernel_types: ClassVar[Mapping[str, type[Kernel]]] = { "fwd_train_whole": BatchNormFwdTrainWholeKernel, @@ -51,7 +51,7 @@ class BatchNormFwdOp(Op): GQA dense decode 的三个实现是第 3 种情况。bs1 服务 batch 为 1 的调用，long-context 服务 KV 较长的调用，二者在 batch 为 1 且 KV 较长时重叠，由 long-context 声明它胜出： ```python -# src/tileops/kernels/attention/gqa/decode.py +# src/tileops/_kernels/attention/gqa/decode.py class GQADecodeLongContextKernel(GQADecodeKernel): general: bool = False preferred_over = frozenset({"gqa_dense_decode_bs1"}) @@ -100,7 +100,7 @@ class GQADecodeLongContextKernel(GQADecodeKernel): 只有一个实现的接口因此只需继承接口。`LayerNormKernel` 只声明了支持的架构与自己的 `entry_for`： ```python -# src/tileops/kernels/norm/layer_norm.py +# src/tileops/_kernels/norm/layer_norm.py class LayerNormKernel(Kernel, LayerNormFwdInterface): supported_archs: list[int] = [80, 86, 89, 90] @@ -123,10 +123,10 @@ class LayerNormKernel(Kernel, LayerNormFwdInterface): | 3 | 某个 shape 范围或某个架构上有更快的算法 | 新实现 | | 4 | 同一算法的 split 数为 1 或大于 1 | 实现内部的参数 | -新接口需要一个 call spec 类型与一个接口类。family 下有多个 kernel 文件时，两者都写在 `src/tileops/kernels//call_spec.py`；只有一个 kernel 文件时，两者都写在这个 kernel 文件中。接口类按 `{Name}{Fwd|Bwd}Interface` 命名，表示变体的词写在方向后缀之前，例如 `BatchNormTrainFwdInterface`，由 pre-commit 的 `interface-names-lint` 检查。同一 family 的多个接口可以共用一个 call spec 类型，例如 BatchNorm 与 InstanceNorm 的六个接口都用 `BatchNormCall`。 +新接口需要一个 call spec 类型与一个接口类。family 下有多个 kernel 文件时，两者都写在 `src/tileops/_kernels//call_spec.py`；只有一个 kernel 文件时，两者都写在这个 kernel 文件中。接口类按 `{Name}{Fwd|Bwd}Interface` 命名，表示变体的词写在方向后缀之前，例如 `BatchNormTrainFwdInterface`，由 pre-commit 的 `interface-names-lint` 检查。同一 family 的多个接口可以共用一个 call spec 类型，例如 BatchNorm 与 InstanceNorm 的六个接口都用 `BatchNormCall`。 ```python -# src/tileops/kernels/norm/call_spec.py +# src/tileops/_kernels/norm/call_spec.py @dataclasses.dataclass(frozen=True) class LayerNormCall(CallSpec): """The facts that select an implementation normalizing trailing rows and build it. @@ -161,7 +161,7 @@ class LayerNormFwdInterface(KernelInterface): ``` - call spec 的字段是各实现在 `refusal`、`entry_for` 中读取的调用事实：shape、dtype、语义参数（包括 op 构造时固定的参数）与 `device`。字段只能是不可变的值。不放张量内容、调优策略与优先级；设备事实由 `CallSpec` 提供。 -- 接口类的 `request` 指向 call spec 类型；抽象方法 `forward` 的参数表是 op 调用 kernel 时传入的参数，docstring 写明每个张量的 shape、dtype、内存布局、设备、是否被原地写入，以及返回值。backend 只依据这份契约编写实现。 +- 接口类的 `request` 指向 call spec 类型；抽象方法 `forward` 的参数表是 op 调用 kernel 时传入的参数，docstring 写明每个张量的 shape、dtype、内存布局、设备、是否被原地写入，以及返回值。实现只依据这份契约编写。 - op 在 `forward` 之外还会在 entry 上调用的方法，一并写成接口的抽象方法。实现必须实现接口的全部抽象方法，否则这个类不能实例化。 - 各实现必须一致的取值，写成接口的普通类方法，由接口算出。例如 op 要在选择之前按大小分配一个输出缓冲区时，这个大小由接口的类方法从 call spec 算出，而不是从已选中的实现上读。 - docstring 写的必须是实现真正做到的：声明输入连续，op 就要在调用前调用 `.contiguous()`；实现不保证的行为（例如输入含 NaN 时的结果）写明未定义，不写成保证。 @@ -174,8 +174,8 @@ class LayerNormFwdInterface(KernelInterface): ```python # tests/ops/test_family_dispatch.py def test_gemm_k_too_narrow_to_vectorize_is_refused_during_selection() -> None: - op = GemmFwdOp() - call = GemmCall(arch=_SM90, sm_count=132, m=64, n=64, k=1, dtype=torch.float16, trans_b=True) + op = GEMMFwdOp() + call = GEMMCall(arch=_SM90, sm_count=132, m=64, n=64, k=1, dtype=torch.float16, trans_b=True) with pytest.raises(ValueError, match="k must span at least one"): op.key_for("gemm", call) diff --git a/docs/user-guide/manifest/calls.md b/docs/user-guide/manifest/calls.md index bff7218d..f1a50c59 100644 --- a/docs/user-guide/manifest/calls.md +++ b/docs/user-guide/manifest/calls.md @@ -56,7 +56,7 @@ The construction phase first solves the indices that construction-time tensors d | No. | spec | Declaration | Solved index | | --- | --- | --- | --- | -| 1 | SiluAndMul | `x: "[M, 2 * N]"` | `N`, also checking that the axis length is even | +| 1 | SiLUAndMul | `input: "[M, 2 * N]"` | `N`, also checking that the axis length is even | | 2 | GQA varlen | `cu_seqlens_q: "[B + 1]"` | `B` | | 3 | GemmW4A16 | `activation: "[M, K]"`, `packed_weight: "[N, K // 2]"` | `K` is solved from `activation`; `K // 2` in `packed_weight` is used only for checking | diff --git a/docs/user-guide/manifest/calls.zh.md b/docs/user-guide/manifest/calls.zh.md index e9fe5c17..59be7800 100644 --- a/docs/user-guide/manifest/calls.zh.md +++ b/docs/user-guide/manifest/calls.zh.md @@ -56,7 +56,7 @@ op 在 `interfaces` 中声明它调用 kernel 的位置，每个位置对应一 | No. | spec | 声明 | 求得的 index | | --- | --- | --- | --- | -| 1 | SiluAndMul | `x: "[M, 2 * N]"` | `N`，同时检查轴长为偶数 | +| 1 | SiLUAndMul | `input: "[M, 2 * N]"` | `N`，同时检查轴长为偶数 | | 2 | GQA varlen | `cu_seqlens_q: "[B + 1]"` | `B` | | 3 | GemmW4A16 | `activation: "[M, K]"`、`packed_weight: "[N, K // 2]"` | `K` 由 `activation` 求得，`packed_weight` 的 `K // 2` 只用于检查 | diff --git a/docs/user-guide/manifest/examples.md b/docs/user-guide/manifest/examples.md index 09de77c3..2c7be3d5 100644 --- a/docs/user-guide/manifest/examples.md +++ b/docs/user-guide/manifest/examples.md @@ -2,11 +2,11 @@ This page collects five real specs from the manifest. Each spec is followed by a table that maps each form used in it to the section that describes it. The specs are taken from `src/tileops/manifest/`, with comments and some workload rows removed. -## 1. GemmFwdOp: type family {#gemm} +## 1. GEMMFwdOp: type family {#gemm} ```yaml # src/tileops/manifest/spec/gemm.yaml -GemmFwdOp: +GEMMFwdOp: ref_api: "torch.matmul" family: gemm status: implemented @@ -38,7 +38,7 @@ GemmFwdOp: flops: "2 * M * N * K" ``` -**Table 1** The forms used in `GemmFwdOp` +**Table 1** The forms used in `GEMMFwdOp` | No. | Form in the spec | Described in | | --- | --- | --- | @@ -124,10 +124,10 @@ MeanPoolingFwdOp: accum_dtype: {type: torch.dtype} inputs: x: {dtype: T, shape: "[B, S, H, D]"} - offsets: {dtype: int32, optional: true, shape: "[NS + 1]", values: "prefix_sum(seq_lens)", requires: ["prefix_offsets(S)"]} - indices: {dtype: int32, optional: "present(offsets)", shape: "[NC, 2]", values: "chunk_indices(seq_lens, chunk_size)"} + cu_seqlens: {dtype: int32, optional: true, shape: "[NS + 1]", values: "prefix_sum(seq_lens)", requires: ["prefix_offsets(S)"]} + chunk_indices: {dtype: int32, optional: "present(cu_seqlens)", shape: "[NC, 2]", values: "chunk_indices(seq_lens, chunk_size)"} outputs: - output: {dtype: T, shape: "PoolOut[present(offsets), B, NC, S, chunk_size, H, D]"} + output: {dtype: T, shape: "PoolOut[present(cu_seqlens), B, NC, S, chunk_size, H, D]"} shape_rules: - "chunk_size > 0 and chunk_size % 32 == 0" - "D <= 128 or D % 128 == 0" @@ -135,29 +135,29 @@ MeanPoolingFwdOp: workloads: - {B: 1, S: 8192, H: 64, D: 128, chunk_size: 64, accum_dtype: float32, dtype_cases: [{T: float16}, {T: bfloat16}], label: uniform-8k} - - {B: 1, S: 8192, H: 64, D: 128, seq_lens: "repeat(2048, 4)", chunk_size: 64, accum_dtype: float32, some: [offsets], dtype_cases: [{T: float16}, {T: float32}], label: ragged-even} + - {B: 1, S: 8192, H: 64, D: 128, seq_lens: "repeat(2048, 4)", chunk_size: 64, accum_dtype: float32, some: [cu_seqlens], dtype_cases: [{T: float16}, {T: float32}], label: ragged-even} # remaining rows omitted roofline: - flops: "B * S * H * D + B * (NC if present(offsets) else ceil_div(S, chunk_size)) * H * D" + flops: "B * S * H * D + B * (NC if present(cu_seqlens) else ceil_div(S, chunk_size)) * H * D" ``` **Table 3** The forms used in `MeanPoolingFwdOp` | No. | Form in the spec | Described in | | --- | --- | --- | -| 1 | `indices` is declared `optional: "present(offsets)"`, so it is passed together with `offsets` or omitted together with it | [Extensions § 1](extensions.md#presence) | -| 2 | the type family `PoolOut` uses `present(offsets)` as its discriminant | [Extensions § 2](extensions.md#type-family) | -| 3 | the first row does not pass `offsets` and does not use `seq_lens`, so the row does not give it; the second row passes `offsets` through `some: [offsets]` and gives `seq_lens`, written here as the value primitive call `"repeat(2048, 4)"` | [Extensions § 1](extensions.md#presence), [Spec fields § 8](writing.md#workloads) | +| 1 | `chunk_indices` is declared `optional: "present(cu_seqlens)"`, so it is passed together with `cu_seqlens` or omitted together with it | [Extensions § 1](extensions.md#presence) | +| 2 | the type family `PoolOut` uses `present(cu_seqlens)` as its discriminant | [Extensions § 2](extensions.md#type-family) | +| 3 | the first row does not pass `cu_seqlens` and does not use `seq_lens`, so the row does not give it; the second row passes `cu_seqlens` through `some: [cu_seqlens]` and gives `seq_lens`, written here as the value primitive call `"repeat(2048, 4)"` | [Extensions § 1](extensions.md#presence), [Spec fields § 8](writing.md#workloads) | | 4 | `NS` and `NC` are solved by unification of the generator results and are not given in the row | [Extensions § 7](extensions.md#generators) | | 5 | `category(accum_dtype) == 'float'` requires the accumulation dtype to be a floating-point type | [Extensions § 9](extensions.md#scalar-dtype) | -| 6 | the roofline formula distinguishes the two cases with `present(offsets)` | [Spec fields § 9](writing.md#roofline) | +| 6 | the roofline formula distinguishes the two cases with `present(cu_seqlens)` | [Spec fields § 9](writing.md#roofline) | -## 4. ReluFwdOp: arbitrary rank and effect {#relu} +## 4. ReLUFwdOp: arbitrary rank and effect {#relu} ```yaml # src/tileops/manifest/spec/elementwise_unary_activation.yaml -ReluFwdOp: +ReLUFwdOp: ref_api: "torch.nn.functional.relu" family: elementwise status: implemented @@ -179,7 +179,7 @@ ReluFwdOp: flops: "prod(S)" ``` -**Table 4** The forms used in `ReluFwdOp` +**Table 4** The forms used in `ReLUFwdOp` | No. | Form in the spec | Described in | | --- | --- | --- | diff --git a/docs/user-guide/manifest/examples.zh.md b/docs/user-guide/manifest/examples.zh.md index 8188c0b5..f113e0d4 100644 --- a/docs/user-guide/manifest/examples.zh.md +++ b/docs/user-guide/manifest/examples.zh.md @@ -2,11 +2,11 @@ 本页收录 manifest 中的五个真实 spec，每个 spec 之后附有一张表，列出其中各处写法对应的说明。这些 spec 摘自 `src/tileops/manifest/`，删去了注释与部分 workload 行。 -## 1. GemmFwdOp：type family {#gemm} +## 1. GEMMFwdOp：type family {#gemm} ```yaml # src/tileops/manifest/spec/gemm.yaml -GemmFwdOp: +GEMMFwdOp: ref_api: "torch.matmul" family: gemm status: implemented @@ -38,7 +38,7 @@ GemmFwdOp: flops: "2 * M * N * K" ``` -**表 1** `GemmFwdOp` 中各处写法的说明 +**表 1** `GEMMFwdOp` 中各处写法的说明 | No. | spec 中的写法 | 说明所在 | | --- | --- | --- | @@ -124,10 +124,10 @@ MeanPoolingFwdOp: accum_dtype: {type: torch.dtype} inputs: x: {dtype: T, shape: "[B, S, H, D]"} - offsets: {dtype: int32, optional: true, shape: "[NS + 1]", values: "prefix_sum(seq_lens)", requires: ["prefix_offsets(S)"]} - indices: {dtype: int32, optional: "present(offsets)", shape: "[NC, 2]", values: "chunk_indices(seq_lens, chunk_size)"} + cu_seqlens: {dtype: int32, optional: true, shape: "[NS + 1]", values: "prefix_sum(seq_lens)", requires: ["prefix_offsets(S)"]} + chunk_indices: {dtype: int32, optional: "present(cu_seqlens)", shape: "[NC, 2]", values: "chunk_indices(seq_lens, chunk_size)"} outputs: - output: {dtype: T, shape: "PoolOut[present(offsets), B, NC, S, chunk_size, H, D]"} + output: {dtype: T, shape: "PoolOut[present(cu_seqlens), B, NC, S, chunk_size, H, D]"} shape_rules: - "chunk_size > 0 and chunk_size % 32 == 0" - "D <= 128 or D % 128 == 0" @@ -135,29 +135,29 @@ MeanPoolingFwdOp: workloads: - {B: 1, S: 8192, H: 64, D: 128, chunk_size: 64, accum_dtype: float32, dtype_cases: [{T: float16}, {T: bfloat16}], label: uniform-8k} - - {B: 1, S: 8192, H: 64, D: 128, seq_lens: "repeat(2048, 4)", chunk_size: 64, accum_dtype: float32, some: [offsets], dtype_cases: [{T: float16}, {T: float32}], label: ragged-even} + - {B: 1, S: 8192, H: 64, D: 128, seq_lens: "repeat(2048, 4)", chunk_size: 64, accum_dtype: float32, some: [cu_seqlens], dtype_cases: [{T: float16}, {T: float32}], label: ragged-even} # 其余行省略 roofline: - flops: "B * S * H * D + B * (NC if present(offsets) else ceil_div(S, chunk_size)) * H * D" + flops: "B * S * H * D + B * (NC if present(cu_seqlens) else ceil_div(S, chunk_size)) * H * D" ``` **表 3** `MeanPoolingFwdOp` 中各处写法的说明 | No. | spec 中的写法 | 说明所在 | | --- | --- | --- | -| 1 | `indices` 声明为 `optional: "present(offsets)"`，因此与 `offsets` 同时传入或同时不传 | [扩展写法 1](extensions.md#presence) | -| 2 | type family `PoolOut` 以 `present(offsets)` 为 discriminant | [扩展写法 2](extensions.md#type-family) | -| 3 | 第一行不传 `offsets`，没有用到 `seq_lens`，因此行中不给出它；第二行通过 `some: [offsets]` 传入 `offsets`，同时给出 `seq_lens`，这里写成值 primitive 调用 `"repeat(2048, 4)"` | [扩展写法 1](extensions.md#presence)、[写一个 spec 8](writing.md#workloads) | +| 1 | `chunk_indices` 声明为 `optional: "present(cu_seqlens)"`，因此与 `cu_seqlens` 同时传入或同时不传 | [扩展写法 1](extensions.md#presence) | +| 2 | type family `PoolOut` 以 `present(cu_seqlens)` 为 discriminant | [扩展写法 2](extensions.md#type-family) | +| 3 | 第一行不传 `cu_seqlens`，没有用到 `seq_lens`，因此行中不给出它；第二行通过 `some: [cu_seqlens]` 传入 `cu_seqlens`，同时给出 `seq_lens`，这里写成值 primitive 调用 `"repeat(2048, 4)"` | [扩展写法 1](extensions.md#presence)、[写一个 spec 8](writing.md#workloads) | | 4 | `NS` 与 `NC` 由 generator 结果的 unification 求得，不在行中给出 | [扩展写法 7](extensions.md#generators) | | 5 | `category(accum_dtype) == 'float'` 限制累加 dtype 必须是浮点类型 | [扩展写法 9](extensions.md#scalar-dtype) | -| 6 | roofline 公式通过 `present(offsets)` 区分两种情况 | [写一个 spec 9](writing.md#roofline) | +| 6 | roofline 公式通过 `present(cu_seqlens)` 区分两种情况 | [写一个 spec 9](writing.md#roofline) | -## 4. ReluFwdOp：任意秩与 effect {#relu} +## 4. ReLUFwdOp：任意秩与 effect {#relu} ```yaml # src/tileops/manifest/spec/elementwise_unary_activation.yaml -ReluFwdOp: +ReLUFwdOp: ref_api: "torch.nn.functional.relu" family: elementwise status: implemented @@ -179,7 +179,7 @@ ReluFwdOp: flops: "prod(S)" ``` -**表 4** `ReluFwdOp` 中各处写法的说明 +**表 4** `ReLUFwdOp` 中各处写法的说明 | No. | spec 中的写法 | 说明所在 | | --- | --- | --- | diff --git a/docs/user-guide/manifest/extensions.md b/docs/user-guide/manifest/extensions.md index c096f6f6..4b993c52 100644 --- a/docs/user-guide/manifest/extensions.md +++ b/docs/user-guide/manifest/extensions.md @@ -76,14 +76,14 @@ adts: MGroupedLayout: sum: contiguous: - python: tileops.ops.moe.contracts.ContiguousLayoutSpec + python: tileops._ops.moe.contracts.ContiguousLayoutSpec fields: - packing: {type: "'tight' | 'aligned'", python: tileops.ops.moe.contracts.ContiguousPacking} - metadata_kind: {type: "'physical_psum' | 'per_row'", python: tileops.ops.moe.contracts.ContiguousMetadata} + packing: {type: "'tight' | 'aligned'", python: tileops._ops.moe.contracts.ContiguousPacking} + metadata_kind: {type: "'physical_psum' | 'per_row'", python: tileops._ops.moe.contracts.ContiguousMetadata} alignment: Dim invariant: "(packing == 'tight') == (alignment == 1) and alignment >= 1" masked: - python: tileops.ops.moe.contracts.MaskedLayoutSpec + python: tileops._ops.moe.contracts.MaskedLayoutSpec fields: {max_m: Dim} ``` @@ -138,7 +138,7 @@ An op without effect declarations only reads its inputs and allocates new tensor | No. | Declaration | Meaning | Example | | --- | --- | --- | --- | -| 1 | `buffer: out` on an output | `forward` gets a parameter `out` after all inputs. When the caller passes `out`, the op writes the result into it and returns it; otherwise the op allocates a new tensor. `out` has the same shape and dtype as that output | `output` of MoEGroupedGemm | +| 1 | `buffer: out` on an output | `forward` gets a parameter `out` after all inputs. When the caller passes `out`, the op writes the result into it and returns it; otherwise the op allocates a new tensor. `out` has the same shape and dtype as that output | `output` of MoEGroupedGEMM | | 2 | `mutated: true` on an input | the op may write this input, and its contents before the call take part in the computation | | | 3 | `mutated: true` and `write_only: true` on an input | a result buffer that must be passed: the op overwrites it, and the result depends only on the other inputs; if the op returns `None`, `outputs` is empty | `output` of FusedMoEExperts | | 4 | `mutated: ""` on an input | the op writes this input only when the expression is true | `mutated: inplace` of activation functions | @@ -171,26 +171,26 @@ Varlen, paged and similar ops use metadata tensors such as `cu_seqlens` and `blo ```yaml forall: {B: Dim, S: Dim, H: Dim, D: Dim, NS: Dim, NC: Dim, T: "DType[...]", seq_lens: "Seq[Int]"} inputs: - offsets: {dtype: int32, optional: true, shape: "[NS + 1]", - values: "prefix_sum(seq_lens)", requires: ["prefix_offsets(S)"]} - indices: {dtype: int32, optional: "present(offsets)", shape: "[NC, 2]", - values: "chunk_indices(seq_lens, chunk_size)"} + cu_seqlens: {dtype: int32, optional: true, shape: "[NS + 1]", + values: "prefix_sum(seq_lens)", requires: ["prefix_offsets(S)"]} + chunk_indices: {dtype: int32, optional: "present(cu_seqlens)", shape: "[NC, 2]", + values: "chunk_indices(seq_lens, chunk_size)"} ``` The rules for generators: - The arguments of a generator are given by the workload row, for example the list of lengths `seq_lens`. A name of kind `Seq[Int]` in `forall` can be used only as a generator argument. A `Seq[Int]` in a workload row can be written as a list of integers or as a call to a value primitive that returns a list, for example `seq_lens: "repeat(512, 64)"`. -- At instantiation, the generator's result is unified with the declared shape, which solves the other indices in the shape, so they are not given in the row. In the example above, `NS` and `NC` are solved from the generated `offsets` and `indices` respectively. In an actual call, these indices are likewise solved from the input tensors. +- At instantiation, the generator's result is unified with the declared shape, which solves the other indices in the shape, so they are not given in the row. In the example above, `NS` and `NC` are solved from the generated `cu_seqlens` and `chunk_indices` respectively. In an actual call, these indices are likewise solved from the input tensors. - A generator is either deterministic or uses private random numbers derived from the workload seed, so the same row generates the same values every time. - The rank of each generator result is fixed, or determined by its shape arguments. - A generated tensor declares an integer dtype (`int32` or `int64`). When an argument is outside its domain, or the result is outside the range of the declared dtype, the generator raises an error. -- A generator argument can be a primitive that returns a list, for example `as_tensor(balanced_sizes(M, G))` in GroupedGemm. +- A generator argument can be a primitive that returns a list, for example `as_tensor(balanced_sizes(M, G))` in GroupedGEMM. The rules for `requires`: - `requires` lists predicates that constrain the contents of a metadata tensor. For example, `prefix_offsets(S)` requires the tensor's first element to be 0, the elements to be non-decreasing, and the last element to be `S`. The contents of the constrained tensor are the predicate's implicit first argument. - Each predicate reads the constrained tensor at a fixed rank; predicates that give elementwise bounds (such as `in_range`) can apply to tensors of any rank. -- A predicate argument can be another metadata tensor, so that one predicate constrains the relation between two tensors. For example, `batch_offsets` in GroupedGemm declares `requires: ["exclusive_prefix_of(batch_sizes)"]`, which requires it to be the exclusive prefix sum of `batch_sizes`. On every branch where the constrained tensor exists, the tensor used as the argument must also exist. +- A predicate argument can be another metadata tensor, so that one predicate constrains the relation between two tensors. For example, `batch_offsets` in GroupedGEMM declares `requires: ["exclusive_prefix_of(batch_sizes)"]`, which requires it to be the exclusive prefix sum of `batch_sizes`. On every branch where the constrained tensor exists, the tensor used as the argument must also exist. - The validator checks `requires` at instantiation against the generated values. In an actual call, the caller guarantees these constraints and the op does not check tensor contents; the validator therefore also checks that the predicate is well defined on every branch where the constrained tensor exists. - A tensor that declares `requires` must also declare `values`. diff --git a/docs/user-guide/manifest/extensions.zh.md b/docs/user-guide/manifest/extensions.zh.md index f6cd05c5..b05dbfb8 100644 --- a/docs/user-guide/manifest/extensions.zh.md +++ b/docs/user-guide/manifest/extensions.zh.md @@ -76,14 +76,14 @@ adts: MGroupedLayout: sum: contiguous: - python: tileops.ops.moe.contracts.ContiguousLayoutSpec + python: tileops._ops.moe.contracts.ContiguousLayoutSpec fields: - packing: {type: "'tight' | 'aligned'", python: tileops.ops.moe.contracts.ContiguousPacking} - metadata_kind: {type: "'physical_psum' | 'per_row'", python: tileops.ops.moe.contracts.ContiguousMetadata} + packing: {type: "'tight' | 'aligned'", python: tileops._ops.moe.contracts.ContiguousPacking} + metadata_kind: {type: "'physical_psum' | 'per_row'", python: tileops._ops.moe.contracts.ContiguousMetadata} alignment: Dim invariant: "(packing == 'tight') == (alignment == 1) and alignment >= 1" masked: - python: tileops.ops.moe.contracts.MaskedLayoutSpec + python: tileops._ops.moe.contracts.MaskedLayoutSpec fields: {max_m: Dim} ``` @@ -138,7 +138,7 @@ primitive 是表达式中可以调用的内建函数，例如 `broadcast`、`red | No. | 声明 | 含义 | 例子 | | --- | --- | --- | --- | -| 1 | 输出上的 `buffer: out` | `forward` 在所有输入之后增加参数 `out`。调用方传入 `out` 时，op 将结果写入并返回它；未传入时，op 分配新张量。`out` 与该输出的形状和 dtype 相同 | MoEGroupedGemm 的 `output` | +| 1 | 输出上的 `buffer: out` | `forward` 在所有输入之后增加参数 `out`。调用方传入 `out` 时，op 将结果写入并返回它；未传入时，op 分配新张量。`out` 与该输出的形状和 dtype 相同 | MoEGroupedGEMM 的 `output` | | 2 | 输入上的 `mutated: true` | op 可能写入这个输入，它在调用前的内容参与计算 | | | 3 | 输入上的 `mutated: true` 与 `write_only: true` | 必须传入的结果缓冲：op 覆盖写入，结果只取决于其他输入；如果 op 返回 `None`，`outputs` 为空 | FusedMoEExperts 的 `output` | | 4 | 输入上的 `mutated: ""` | 仅当表达式为 true 时，op 才写入这个输入 | 激活函数的 `mutated: inplace` | @@ -171,26 +171,26 @@ varlen、paged 等 op 使用 metadata 张量，例如 `cu_seqlens`、`block_tabl ```yaml forall: {B: Dim, S: Dim, H: Dim, D: Dim, NS: Dim, NC: Dim, T: "DType[...]", seq_lens: "Seq[Int]"} inputs: - offsets: {dtype: int32, optional: true, shape: "[NS + 1]", - values: "prefix_sum(seq_lens)", requires: ["prefix_offsets(S)"]} - indices: {dtype: int32, optional: "present(offsets)", shape: "[NC, 2]", - values: "chunk_indices(seq_lens, chunk_size)"} + cu_seqlens: {dtype: int32, optional: true, shape: "[NS + 1]", + values: "prefix_sum(seq_lens)", requires: ["prefix_offsets(S)"]} + chunk_indices: {dtype: int32, optional: "present(cu_seqlens)", shape: "[NC, 2]", + values: "chunk_indices(seq_lens, chunk_size)"} ``` generator 的规则如下： - generator 的实参由 workload 行给出，例如长度列表 `seq_lens`。`forall` 中 kind 为 `Seq[Int]` 的名字只能用作 generator 的实参。workload 行中的 `Seq[Int]` 可以写成整数列表，也可以写成返回列表的值 primitive 调用，例如 `seq_lens: "repeat(512, 64)"`。 -- 实例化时，generator 的结果与声明的形状做 unification，形状中的其他 index 由此求得，不需要在行中给出。上例中的 `NS` 与 `NC` 分别由 `offsets` 与 `indices` 的生成结果求得。实际调用时，这些 index 同样由输入张量求得。 +- 实例化时，generator 的结果与声明的形状做 unification，形状中的其他 index 由此求得，不需要在行中给出。上例中的 `NS` 与 `NC` 分别由 `cu_seqlens` 与 `chunk_indices` 的生成结果求得。实际调用时，这些 index 同样由输入张量求得。 - generator 或者是确定性的，或者使用由 workload 种子派生的私有随机数，因此同一条行每次生成的取值都相同。 - 每个 generator 结果的秩是固定的，或者由其形状参数决定。 - 生成的张量声明整数 dtype（`int32` 或 `int64`）；参数超出定义域，或结果超出所声明 dtype 的范围时，generator 会报错。 -- generator 的实参可以是返回列表的 primitive，例如 GroupedGemm 的 `as_tensor(balanced_sizes(M, G))`。 +- generator 的实参可以是返回列表的 primitive，例如 GroupedGEMM 的 `as_tensor(balanced_sizes(M, G))`。 `requires` 的规则如下： - `requires` 列出约束 metadata 张量内容的谓词。例如 `prefix_offsets(S)` 要求张量首项为 0、单调不减、末项为 `S`。被约束张量的内容是谓词隐含的第一个实参。 - 每个谓词按固定的秩读取被约束的张量；逐元素给出上下界的谓词（如 `in_range`）可以作用于任意秩的张量。 -- 谓词的实参可以是另一个 metadata 张量，这样一条谓词就能约束两个张量之间的关系，例如 GroupedGemm 的 `batch_offsets` 声明 `requires: ["exclusive_prefix_of(batch_sizes)"]`，要求它是 `batch_sizes` 的 exclusive 前缀和。在被约束的张量存在的每个分支上，作为实参的那个张量也必须存在。 +- 谓词的实参可以是另一个 metadata 张量，这样一条谓词就能约束两个张量之间的关系，例如 GroupedGEMM 的 `batch_offsets` 声明 `requires: ["exclusive_prefix_of(batch_sizes)"]`，要求它是 `batch_sizes` 的 exclusive 前缀和。在被约束的张量存在的每个分支上，作为实参的那个张量也必须存在。 - validator 在实例化时用生成的取值检查 `requires`。在实际调用中，这些约束由调用方保证，op 不检查张量的内容；因此 validator 还会检查，在被约束的张量存在的每个分支上，谓词都是良定义的。 - 声明了 `requires` 的张量也必须声明 `values`。 diff --git a/docs/user-guide/manifest/index.md b/docs/user-guide/manifest/index.md index e17c3b84..b0668422 100644 --- a/docs/user-guide/manifest/index.md +++ b/docs/user-guide/manifest/index.md @@ -36,18 +36,18 @@ In the figure below, purple components are written by the developer, cyan compon | No. | Part | Location | Contents | | --- | --- | --- | --- | | 1 | spec | `src/tileops/manifest/spec/.yaml` | signature, workload rows, roofline formulas | -| 2 | kernel interface and implementations | `src/tileops/kernels/`; the kernel interface and the call spec usually live in the family's `call_spec.py`, and a family with a single kernel file writes them in that file | The kernel interface fixes the type of the call spec and the parameters of `forward`; an implementation is a subclass of `Kernel` that also inherits the corresponding kernel interface. The default `entry_for(call)` constructs the implementation directly from the call spec; it is overridden only when that does not express the construction arguments correctly | -| 3 | Op class | `src/tileops/ops/`, exported by `tileops.` | an `__init__` that matches `params` and calls `super().__init__(target=target)` once the manifest parameters are assigned; `kernel_types` (key to implementation class) and `interfaces` (call site to kernel interface); the computation `forward`, which constructs the call spec and obtains the entry to call through `kernel_for`; a Google-style docstring, from which the docs site generates the API reference | +| 2 | kernel interface and implementations | `src/tileops/_kernels/`; the kernel interface and the call spec usually live in the family's `call_spec.py`, and a family with a single kernel file writes them in that file | The kernel interface fixes the type of the call spec and the parameters of `forward`; an implementation is a subclass of `Kernel` that also inherits the corresponding kernel interface. The default `entry_for(call)` constructs the implementation directly from the call spec; it is overridden only when that does not express the construction arguments correctly | +| 3 | Op class | `src/tileops/_ops/`, exported by `tileops.` | an `__init__` that matches `params` and calls `super().__init__(target=target)` once the manifest parameters are assigned; `kernel_types` (key to implementation class) and `interfaces` (call site to kernel interface); the computation `forward`, which constructs the call spec and obtains the entry to call through `kernel_for`; a Google-style docstring, from which the docs site generates the API reference | | 4 | reference implementation | `workloads/` | `ref_program` on the workload class named after the op, or on a parameterized workload class shared by the family, plus the input construction that the workload rows cannot determine | | 5 | correctness tests | `tests/ops/` | numerical tolerances, and the shapes needed to cover every branch of the kernel | | 6 | benchmark function | the benchmark file in `benchmarks/ops/` for the module the op belongs to; a new file is added only when no suitable one exists | a test function parameterized with `bench.cases(Op)`, and the choice of implementations to compare against, such as the torch reference implementation and kernels from other libraries; the op also needs a case entry in `benchmarks/_cases/` | The following are written only when needed: -- when an inline formula cannot express the cost, a roofline `func` in `tileops.perf.formulas`; +- when an inline formula cannot express the cost, a roofline `func` in `tileops._perf.formulas`; - when the op's FLOPs are a matrix-multiply contraction and the best implementation should use tensor cores, an override of `roof_key()`. It states the compute unit the best implementation should use, independent of the unit the current kernel actually uses; - when one kernel interface has several implementations, a declaration on each implementation of which calls it serves (`refusal`) and which devices it runs on (`devices`, `supported_archs`); when the applicability of implementations overlaps, `preferred_over` declares which one takes precedence, and the fallback implementation declares `general = True`, see [Adding a kernel to an op](../dispatch/writing.md#rule); -- an external backend does not modify the op. Before an op instance is constructed, it calls `tileops.backend.register_kernel_type("", "", implementation class)` to add an implementation for a kernel interface. The first argument is the op name in the manifest, and `key` must be a key the op does not use yet. The new implementation enters only instances constructed afterwards, and it is checked against the kernel interface when the instance is constructed, see [How a backend joins TileOPs](../dispatch/backends.md#register); +- an external backend does not modify the op. It calls `tileops.backend.register_kernel_builder` to register a target's builder for the op, and that target takes over every call of the op. Its `op` argument is the op name in the manifest, see [How a backend joins TileOPs](../dispatch/backends.md#target); - a composite op declares `delegate_types` on the class, matching the spec's `composition`; it declares `kernel_types` and `interfaces` only when it also runs kernels of its own; - an op whose spec has a call-time tensor input and no `composition` gets a compile boundary generated when the class is defined; it must pass a cold `fullgraph=True` compile, with a cold-compile test registered in `tests/compile_contract.py`. @@ -77,7 +77,7 @@ An op call has two phases: construction and call. In the figure below, cyan part - In the call phase, the generated checks first determine the inputs, select the branch and infer the indices, and only then call the implementation written by the developer. After the implementation returns, they check the number, shape, dtype, device and memory layout of the outputs, and verify that `out` and `alias` outputs are the corresponding tensor objects. The implementation therefore does not repeat checks the spec already declares. - The implementation obtains the entry to call through `kernel_for(interface, call)`; the entry is usually a single kernel. The Op base class looks up each distinct call spec only once. When a call spec first appears, the Op base class selects exactly one of the kernel interface's implementations, whose `entry_for(call)` returns the build identity and the factory. Within one kernel interface, an entry with the same implementation class and the same build identity is constructed only once. The selection rules are in [Calls and validation § 2](calls.md#selection). - For an op with a compile boundary, `Op.__call__` calls the generated `_call_boundary`, which runs `forward` inside the node through the operator; for an op without one, `Op.__call__` runs `forward` directly. Either way, the computation is written in `forward`. -- When a target has registered a builder for the op through `register_kernel_builder`, the kernel returned by the target is called after the generated checks, and the op's `forward` does not run. Implementations added through `register_kernel_type` belong to the in-tree path and are still selected by `forward` through `kernel_for`. +- When a target has registered a builder for the op through `register_kernel_builder`, the kernel returned by the target is called after the generated checks, and the op's `forward` does not run. - If every tensor a call writes (each output and each written input) has no elements, neither the in-tree implementation nor the target runs: new outputs are created on the call device with the checked shapes and dtypes, and `out` and the written inputs are returned unchanged. When the inputs are empty but the outputs are not, the implementation runs as usual. The details of each check are in [Calls and validation § 1](calls.md#call). @@ -89,7 +89,7 @@ The left side of the figure below shows the developer's steps; the right side sh ![Workflow for adding an op](img/new-op.svg) 1. The spec is committed with `status: spec-only`. At this point the validator runs only static checks and does not require the code to exist. -2. The developer writes the kernel and the op class, and exports the op from the package that holds the implementation (for example `src/tileops/ops/reduction/__init__.py`) and from the public module `src/tileops/.py`; the `__all__` of both must include it. As soon as the op class exists, the methods generated from the signature are added to it. +2. The developer writes the kernel and the op class, and exports the op from the package that holds the implementation (for example `src/tileops/_ops/reduction/__init__.py`) and from the public module `src/tileops/.py`; the `__all__` of both must include it. As soon as the op class exists, the methods generated from the signature are added to it. 3. The developer writes the reference implementation, the correctness tests and the benchmark function, adding them to the family's or module's existing file if there is one. See the design document [Testing](../../design/testing.md). - Correctness tests use the `BUILTIN` target by default and take their device from `workloads.device.run_device()`. Device availability is checked with `workloads.device.run_device_available()`, not `torch.cuda.is_available()`. - A test that depends on in-tree kernel state is marked `pytest.mark.in_tree_kernels` or passes `target=BUILTIN` explicitly. diff --git a/docs/user-guide/manifest/index.zh.md b/docs/user-guide/manifest/index.zh.md index 93485055..d6c8b687 100644 --- a/docs/user-guide/manifest/index.zh.md +++ b/docs/user-guide/manifest/index.zh.md @@ -36,18 +36,18 @@ TileOPs 的组件分为四层： | No. | 部分 | 位置 | 内容 | | --- | --- | --- | --- | | 1 | spec | `src/tileops/manifest/spec/.yaml` | 签名、workload 行、roofline 公式 | -| 2 | kernel 接口与实现 | `src/tileops/kernels/`；kernel 接口与 call spec 通常放在 family 的 `call_spec.py` 中，只有一个 kernel 文件的 family 写在该文件里 | kernel 接口规定 call spec 的类型与 `forward` 的参数；实现是 `Kernel` 的子类，并继承相应的 kernel 接口。默认的 `entry_for(call)` 直接用 call spec 构造实现，不能准确表示构造参数时再覆写它 | -| 3 | Op 类 | `src/tileops/ops/`，并由 `tileops.` 导出 | 与 `params` 一致的 `__init__`，赋值 manifest 参数后调用 `super().__init__(target=target)`；`kernel_types`（键到实现类）与 `interfaces`（调用位置到 kernel 接口）；计算本体 `forward`，在其中构造 call spec，并通过 `kernel_for` 取得要调用的 entry；Google 风格的 docstring，文档站的 API 参考由它生成 | +| 2 | kernel 接口与实现 | `src/tileops/_kernels/`；kernel 接口与 call spec 通常放在 family 的 `call_spec.py` 中，只有一个 kernel 文件的 family 写在该文件里 | kernel 接口规定 call spec 的类型与 `forward` 的参数；实现是 `Kernel` 的子类，并继承相应的 kernel 接口。默认的 `entry_for(call)` 直接用 call spec 构造实现，不能准确表示构造参数时再覆写它 | +| 3 | Op 类 | `src/tileops/_ops/`，并由 `tileops.` 导出 | 与 `params` 一致的 `__init__`，赋值 manifest 参数后调用 `super().__init__(target=target)`；`kernel_types`（键到实现类）与 `interfaces`（调用位置到 kernel 接口）；计算本体 `forward`，在其中构造 call spec，并通过 `kernel_for` 取得要调用的 entry；Google 风格的 docstring，文档站的 API 参考由它生成 | | 4 | 参考实现 | `workloads/` | 以该 op 命名的 workload 类，或 family 共用的参数化 workload 类上的 `ref_program`，以及 workload 行无法确定的输入构造 | | 5 | 正确性测试 | `tests/ops/` | 数值容差，以及覆盖 kernel 各分支所需的形状 | | 6 | benchmark 函数 | `benchmarks/ops/` 中该 op 所属模块的 benchmark 文件，没有合适的文件时再新增 | 以 `bench.cases(Op)` 参数化测试函数，并选择对比实现，例如 torch 参考实现与其他库的 kernel；该 op 还需要在 `benchmarks/_cases/` 中有一个 case 注册项 | 以下几项只在需要时编写： -- 内联公式无法表达代价时，在 `tileops.perf.formulas` 中编写 roofline `func`； +- 内联公式无法表达代价时，在 `tileops._perf.formulas` 中编写 roofline `func`； - op 的 FLOPs 属于矩阵乘收缩、最优实现应使用 tensor core 时，覆写 `roof_key()`。它表示最优实现应当使用的计算单元，与当前 kernel 实际使用的单元无关； - 同一个 kernel 接口有多个实现时，在各实现上声明它服务哪些调用（`refusal`）与能在哪些设备上运行（`devices`、`supported_archs`）；实现之间的适用范围有重叠时，用 `preferred_over` 声明哪个优先，兜底的实现声明 `general = True`，见[如何为 op 新增 kernel](../dispatch/writing.md#rule)； -- 外部 backend 不修改 op，而是在构造 op 实例之前调用 `tileops.backend.register_kernel_type("", "", 实现类)`，为某个 kernel 接口增加实现。第一个参数是 manifest 中的 op 名，`key` 必须是该 op 尚未使用的新键；新增的实现只进入之后构造的实例，并在实例构造时接受 kernel 接口检查，见 [backend 如何接入 TileOPs](../dispatch/backends.md#register)； +- 外部 backend 不修改 op，而是调用 `tileops.backend.register_kernel_builder`，为 op 注册一个 target 的 builder，由这个 target 接管 op 的全部调用。参数 `op` 是 manifest 中的 op 名，见 [backend 如何接入 TileOPs](../dispatch/backends.md#target)； - 复合 op 在类上声明 `delegate_types`，与 spec 的 `composition` 对应；复合 op 也运行自己的 kernel 时，才声明 `kernel_types` 与 `interfaces`； - spec 有调用期张量输入、且没有 `composition` 的 op 在类定义时自动生成编译边界，必须通过冷启动的 `fullgraph=True` 编译，并在 `tests/compile_contract.py` 中登记冷编译测试。 @@ -77,7 +77,7 @@ manifest 测试与 benchmark 所用的调用，包括形状、dtype 与输入张 - 调用阶段，生成的检查先确定输入、选择分支并推断 index，然后才调用开发者编写的实现；实现返回之后，再检查输出的数量、形状、dtype、设备与内存布局，并核对 `out` 与 `alias` 输出是否就是对应的张量对象。因此实现中不需要重复 spec 已经声明的检查。 - 实现通过 `kernel_for(interface, call)` 取得要调用的 entry，通常就是一个 kernel。同一个 call spec 再次出现时，Op 基类只做一次查找；首次出现时，Op 基类在该 kernel 接口的实现中选出唯一一个，由它的 `entry_for(call)` 返回 build identity 与构建函数，在同一个 kernel 接口内，同一实现类、相同 build identity 的 entry 只构造一次。选择规则见[调用与校验 2](calls.md#selection)。 - 有编译边界的 op，`Op.__call__` 调用生成的 `_call_boundary`，经 operator 在节点内部运行 `forward`；没有编译边界的 op，`Op.__call__` 直接运行 `forward`。两种情况下计算本体都写在 `forward` 中。 -- 某个 target 通过 `register_kernel_builder` 为 op 注册了 builder 时，生成的检查之后调用 target 返回的 kernel，op 的 `forward` 不执行。通过 `register_kernel_type` 增加的实现属于 in-tree 路径，仍由 `forward` 经 `kernel_for` 选中。 +- 某个 target 通过 `register_kernel_builder` 为 op 注册了 builder 时，生成的检查之后调用 target 返回的 kernel，op 的 `forward` 不执行。 - 如果一次调用写入的所有张量（各输出与被写入的输入）都不含元素，in-tree 实现与 target 都不执行：新的输出按检查过的形状与 dtype 在调用设备上创建，`out` 与被写入的输入原样返回。输入为空而输出不为空时，照常执行实现。 各检查的细节见[调用与校验 1](calls.md#call)。 @@ -89,7 +89,7 @@ manifest 测试与 benchmark 所用的调用，包括形状、dtype 与输入张 ![新增一个 op 的流程](img/new-op.svg) 1. spec 以 `status: spec-only` 提交。此时 validator 只做静态检查，不要求代码存在。 -2. 开发者编写 kernel 与 op 类，并在实现所在的包（例如 `src/tileops/ops/reduction/__init__.py`）与公开模块 `src/tileops/.py` 中导出 op，两处的 `__all__` 都要包含它。op 类一旦存在，由签名生成的方法就会加入这个类。 +2. 开发者编写 kernel 与 op 类，并在实现所在的包（例如 `src/tileops/_ops/reduction/__init__.py`）与公开模块 `src/tileops/.py` 中导出 op，两处的 `__all__` 都要包含它。op 类一旦存在，由签名生成的方法就会加入这个类。 3. 开发者编写参考实现、正确性测试与 benchmark 函数。family 或模块已有对应文件时，在已有文件中添加。详见设计文档 [Testing](../../design/testing.md)。 - 正确性测试默认使用 `BUILTIN` target，测试设备取自 `workloads.device.run_device()`。判断设备是否可用时调用 `workloads.device.run_device_available()`，不调用 `torch.cuda.is_available()`。 - 依赖 in-tree kernel 状态的测试标记 `pytest.mark.in_tree_kernels`，或显式传入 `target=BUILTIN`。 diff --git a/docs/user-guide/manifest/writing.md b/docs/user-guide/manifest/writing.md index 62e16b78..7df2d54a 100644 --- a/docs/user-guide/manifest/writing.md +++ b/docs/user-guide/manifest/writing.md @@ -13,16 +13,16 @@ How to write optional inputs, shapes that vary with parameters, side effects, me ## 1. Structure of a spec {#anatomy} -The SiluAndMul spec uses only the fields above: +The SiLUAndMul spec uses only the fields above: ```yaml -SiluAndMulFwdOp: +SiLUAndMulFwdOp: family: elementwise status: implemented signature: forall: {M: Dim, N: Dim, T: "DType[float16 | bfloat16 | float32]"} inputs: - x: {dtype: T, shape: "[M, 2 * N]"} + input: {dtype: T, shape: "[M, 2 * N]"} outputs: output: {dtype: T, shape: "[M, N]"} workloads: @@ -187,7 +187,7 @@ roofline: flops: "2 * M * N * K" # inline formula; when bytes is omitted it is derived from the signature # or roofline: - func: "tileops.perf.formulas.gqa_fwd_roofline" + func: "tileops._perf.formulas.gqa_dense_fwd_roofline" ``` - Inline formulas are written in the expression language and can refer to the signature's indices, construction parameters and `let` entries, as well as `present(t)`, `bytes(t)` (the byte count of tensor `t`) and built-in primitives. When the cost depends on whether a tensor is passed, the formula can distinguish the cases with `present(...)`. @@ -202,7 +202,7 @@ roofline: - recurrent ops such as linear attention and state-space scans count by the chunked algorithm at the chunk size when the signature has a chunk size, and by per-token recurrence otherwise; - elementwise operations, attention and MoE are counted by the conventions in the design document [Roofline § 1.3](../../design/roofline.md#13-convention); - a path that is an identity map on the call's dtype counts as 0. -- `func` is used when the formula needs Python logic. It points to a module-level function `f(call) -> tuple[int, int]` in `tileops.perf.formulas`. The function reads only its argument `call` and does not read the op instance. `call` is one checked call and provides the following interface: +- `func` is used when the formula needs Python logic. It points to a module-level function `f(call) -> tuple[int, int]` in `tileops._perf.formulas`. The function reads only its argument `call` and does not read the op instance. `call` is one checked call and provides the following interface: **Table 4** The argument `call` of a roofline `func` diff --git a/docs/user-guide/manifest/writing.zh.md b/docs/user-guide/manifest/writing.zh.md index 28114464..60b9a7df 100644 --- a/docs/user-guide/manifest/writing.zh.md +++ b/docs/user-guide/manifest/writing.zh.md @@ -13,16 +13,16 @@ ## 1. spec 的结构 {#anatomy} -SiluAndMul 的 spec 只用到了上述字段： +SiLUAndMul 的 spec 只用到了上述字段： ```yaml -SiluAndMulFwdOp: +SiLUAndMulFwdOp: family: elementwise status: implemented signature: forall: {M: Dim, N: Dim, T: "DType[float16 | bfloat16 | float32]"} inputs: - x: {dtype: T, shape: "[M, 2 * N]"} + input: {dtype: T, shape: "[M, 2 * N]"} outputs: output: {dtype: T, shape: "[M, N]"} workloads: @@ -187,7 +187,7 @@ roofline: flops: "2 * M * N * K" # 内联公式；省略 bytes 时由签名推导 # 或 roofline: - func: "tileops.perf.formulas.gqa_fwd_roofline" + func: "tileops._perf.formulas.gqa_dense_fwd_roofline" ``` - 内联公式使用表达式语言书写，可以引用签名中的 index、构造参数与 `let`，以及 `present(t)`、`bytes(t)`（张量 `t` 的字节数）和内建 primitive。如果代价取决于某个张量是否传入，公式可以使用 `present(...)` 区分。 @@ -198,7 +198,7 @@ roofline: - 字节数表示算法的最少访存量：中间结果不计入，算法只读取一部分的输入按实际读取的不同元素计数；两者在每条 workload 行上相差不到 1% 时，可以按整个张量计。 - 如果推导出的字节数与上述最少访存量不符，spec 需要显式写出 `bytes`，并附带相应的测试。 - `flops` 是算法完成本次调用所需的最少算术量，不是硬件实际执行的指令数：已经算出的值按复用计；linear attention、状态空间扫描等循环 op，签名中带 chunk size 时按该 chunk size 的分块算法计，否则按逐 token 递推计；逐元素运算、attention 与 MoE 按设计文档 [Roofline 1.3](../../design/roofline.md#13-convention) 的约定计数；在本次调用的 dtype 上等同于恒等映射的路径计 0。 -- 公式需要 Python 逻辑时使用 `func`。它指向 `tileops.perf.formulas` 中的一个模块级函数 `f(call) -> tuple[int, int]`。函数只读取参数 `call`，不读取 op 实例。`call` 是一次经过检查的调用，提供以下接口： +- 公式需要 Python 逻辑时使用 `func`。它指向 `tileops._perf.formulas` 中的一个模块级函数 `f(call) -> tuple[int, int]`。函数只读取参数 `call`，不读取 op 实例。`call` 是一次经过检查的调用，提供以下接口： **表 4** roofline `func` 的参数 `call` diff --git a/docs/user-guide/op/composite.md b/docs/user-guide/op/composite.md index 8eb050fd..74f52dec 100644 --- a/docs/user-guide/op/composite.md +++ b/docs/user-guide/op/composite.md @@ -14,8 +14,8 @@ This page describes: ```python class MoEExpertMLPFwdOp(Op): delegate_types: ClassVar[Mapping[str, type[Op]]] = { - "gate_up": MoEGroupedGemmFwdOp, - "down": MoEGroupedGemmFwdOp, + "gate_up": MoEGroupedGEMMFwdOp, + "down": MoEGroupedGEMMFwdOp, } def __init__(self, layout, activation="silu_and_mul", *, target=None): @@ -28,7 +28,7 @@ class MoEExpertMLPFwdOp(Op): ``` The example leaves out the class docstring and `forward`; the full code is in -[`src/tileops/ops/moe/staged.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops/moe/staged.py). +[`src/tileops/_ops/moe/staged.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/_ops/moe/staged.py). Sub-ops follow the same rule as kernels: declared on the class first, then obtained through a method of the base class. @@ -53,7 +53,7 @@ method of the base class. - A repeated call with the same `(stage, identity)` returns the instance already held. - Registering an instance already held under another `(stage, identity)` makes `delegate_for` raise. -- **Derived enumeration.** `held_delegates()`, `iter_kernels()` and `request_tune()` are +- **Derived enumeration.** `_held_delegates()`, `iter_kernels()` and `request_tune()` are derived from the sub-ops `delegate_for` holds; a composite op does not override them. ## 2. Sub-op calls are filed under stages {#stages} diff --git a/docs/user-guide/op/composite.zh.md b/docs/user-guide/op/composite.zh.md index 92934044..dd3a965d 100644 --- a/docs/user-guide/op/composite.zh.md +++ b/docs/user-guide/op/composite.zh.md @@ -13,8 +13,8 @@ ```python class MoEExpertMLPFwdOp(Op): delegate_types: ClassVar[Mapping[str, type[Op]]] = { - "gate_up": MoEGroupedGemmFwdOp, - "down": MoEGroupedGemmFwdOp, + "gate_up": MoEGroupedGEMMFwdOp, + "down": MoEGroupedGEMMFwdOp, } def __init__(self, layout, activation="silu_and_mul", *, target=None): @@ -26,7 +26,7 @@ class MoEExpertMLPFwdOp(Op): self.down = self.delegate_for("down", None, layout=layout) ``` -示例省略了类 docstring 和 `forward`，完整代码见 [`src/tileops/ops/moe/staged.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops/moe/staged.py)。子 op 遵循与 kernel 相同的规则：先在类上声明，再经基类的方法取得。 +示例省略了类 docstring 和 `forward`，完整代码见 [`src/tileops/_ops/moe/staged.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/_ops/moe/staged.py)。子 op 遵循与 kernel 相同的规则：先在类上声明，再经基类的方法取得。 - **声明。** `delegate_types` 按 **stage** 名映射组合 op 可能持有的子 op 类；validator 检查它与 manifest 条目的 composition 一致。 - 映射的顺序就是 stage 的顺序，与 manifest 的 composition 一致。 @@ -42,7 +42,7 @@ class MoEExpertMLPFwdOp(Op): - **一个实例只登记一次。** - 同一个 `(stage, identity)` 的重复调用，返回已持有的实例。 - 把已经持有的实例用另一个 `(stage, identity)` 再次登记，`delegate_for` 报错。 -- **派生的枚举。** `held_delegates()`、`iter_kernels()` 和 `request_tune()` 从 `delegate_for` 持有的子 op 派生，组合 op 不覆盖它们。 +- **派生的枚举。** `_held_delegates()`、`iter_kernels()` 和 `request_tune()` 从 `delegate_for` 持有的子 op 派生，组合 op 不覆盖它们。 ## 2. 子 op 的调用归入 stage {#stages} diff --git a/docs/user-guide/op/img/structure.svg b/docs/user-guide/op/img/structure.svg index e31a03bf..6c986528 100644 --- a/docs/user-guide/op/img/structure.svg +++ b/docs/user-guide/op/img/structure.svg @@ -1 +1 @@ -Op class structureOpkernel_types : key → implementation classinterfaces : interface → interface classdelegate_types : stage → Op classinjected_parameterscompile_op_names_signature : _SignaturePlan (shared by the class's instances) Class attributes (one per op class)  Instance state manifest params (construction-time tensors aside)target_installed_kernel_types : key → implementation class_preferred_over_keys_by_interface_construction_indices_instance_keyconstruction_target_builder_serving_target_entries_by_call : (interface, call) → entry_built_entries : interface → entries_target_kernels : (device, each input's dtype, shape) → kernelbindingconstruction-time tensors_tune_requested_delegates : stage → identity → Op_delegate_stages_effect_branches_last_call : SignatureCall_tune_warnedrunning Methods __call__(*args) — the parameters are forward'sserving_target «property»last_call «property»eval_roofline() «generated»request_tune()kernel_config()iter_kernels()built_entries(interface)held_delegates()external interface__init__(*, target=None)forward(...)kernel_for(interface, call)key_for(interface, call)delegate_for(stage, identity, given=None, /, **params)roof_key()internal interface__init_subclass__() — runs at class definition, installs the generated methods_run_call(inputs, body)_bind_target()_reset_binding()_get_or_build_entry(interface, call)_check_signature(tensors)_infer_output_shapes(...) «generated»implementationRMSNormFwdOpkernel_types : 3 implementationsinterfaces : rms_normforward(x, weight=None)MoEExpertMLPFwdOpdelegate_types = {gate_up, down}forward(expert_input, ...)_CompileBoundaryplan : _SignaturePlanbranches : effect branch → custom opcall(op, inputs, writes, execution)build_call_boundary(cls) : the subclass's _call_boundary()_SignaturePlansigconstructionschecksshapeseffectsroofscheck_construction(op) : indicescheck(op, tensors) : SignatureCalloutput_shapes(op, shapes, dtypes)effect_branch(op, present)eval_roofline(call)«dataclass, frozen»SignatureCallindicesbranch_keydevicetensorstrafficwrittenhas_outmetadatastages : stage → SignatureCall*metadata_values(name)Kerneldevicessupported_archsgeneralpreferred_overrefusal(call) [classmethod]entry_for(call) [classmethod]forward(...)KernelInterfacerequest : CallSpecforward(...)RMSNormFwdInterfacerequest = LayerNormCallforward(x, weight)entryA runtime object: on a cache miss,built by the builder the implementation class's entry_for returnsone kernel or several kernels togetherhere, one RMSNormOnChipKernel instance«Op subclass»MoEGroupedGemmFwdOpforward(...)RMSNormOnChipKernelsupported_archs = [90]preferred_over = {rms_norm_streaming}refusal(call) [classmethod]20..1last_call0..1table[key]0..* +Op class structureOpkernel_types : key → implementation classinterfaces : interface → interface classdelegate_types : stage → Op classinjected_parameterscompile_op_names_signature : _SignaturePlan (shared by the class's instances) Class attributes (one per op class)  Instance state manifest params (construction-time tensors aside)target_installed_kernel_types : key → implementation class_preferred_over_keys_by_interface_construction_indices_instance_keyconstruction_target_builder_serving_target_entries_by_call : (interface, call) → entry_built_entries : interface → entries_target_kernels : (device, each input's dtype, shape) → kernelbindingconstruction-time tensors_tune_requested_delegates : stage → identity → Op_delegate_stages_effect_branches_last_call : SignatureCall_tune_warnedrunning Methods __call__(*args) — the parameters are forward'sserving_target «property»last_call «property»eval_roofline() «generated»request_tune()kernel_config()iter_kernels()_entries_built(interface)_held_delegates()external interface__init__(*, target=None)forward(...)kernel_for(interface, call)key_for(interface, call)delegate_for(stage, identity, given=None, /, **params)roof_key()internal interface__init_subclass__() — runs at class definition, installs the generated methods_run_call(inputs, body)_bind_target()_reset_binding()_get_or_build_entry(interface, call)_check_signature(tensors)_infer_output_shapes(...) «generated»implementationRMSNormFwdOpkernel_types : 3 implementationsinterfaces : rms_normforward(x, weight=None)MoEExpertMLPFwdOpdelegate_types = {gate_up, down}forward(expert_input, ...)_CompileBoundaryplan : _SignaturePlanbranches : effect branch → custom opcall(op, inputs, writes, execution)build_call_boundary(cls) : the subclass's _call_boundary()_SignaturePlansigconstructionschecksshapeseffectsroofscheck_construction(op) : indicescheck(op, tensors) : SignatureCalloutput_shapes(op, shapes, dtypes)effect_branch(op, present)eval_roofline(call)«dataclass, frozen»SignatureCallindicesbranch_keydevicetensorstrafficwrittenhas_outmetadatastages : stage → SignatureCall*metadata_values(name)Kerneldevicessupported_archsgeneralpreferred_overrefusal(call) [classmethod]entry_for(call) [classmethod]forward(...)KernelInterfacerequest : CallSpecforward(...)RMSNormFwdInterfacerequest = LayerNormCallforward(x, weight)entryA runtime object: on a cache miss,built by the builder the implementation class's entry_for returnsone kernel or several kernels togetherhere, one RMSNormOnChipKernel instance«Op subclass»MoEGroupedGEMMFwdOpforward(...)RMSNormOnChipKernelsupported_archs = [90]preferred_over = {rms_norm_streaming}refusal(call) [classmethod]20..1last_call0..1table[key]0..* diff --git a/docs/user-guide/op/img/structure.zh.svg b/docs/user-guide/op/img/structure.zh.svg index e469cb4d..3aecbf2f 100644 --- a/docs/user-guide/op/img/structure.zh.svg +++ b/docs/user-guide/op/img/structure.zh.svg @@ -1 +1 @@ -Op 类结构Opkernel_types : key → 实现类interfaces : interface → 接口类delegate_types : stage → Op 类injected_parameterscompile_op_names_signature : _SignaturePlan（同类实例共享） 类属性（每个 op 类一份）  实例状态 manifest 参数（构造期张量除外）target_installed_kernel_types : key → 实现类_preferred_over_keys_by_interface_construction_indices_instance_key构造_target_builder_serving_target_entries_by_call : (interface, call) → entry_built_entries : interface → entries_target_kernels : (设备, 各输入 dtype、shape) → kernel绑定构造期张量_tune_requested_delegates : stage → identity → Op_delegate_stages_effect_branches_last_call : SignatureCall_tune_warned运行 方法 __call__(*args) —— 参数即 forward 的签名serving_target «property»last_call «property»eval_roofline() «生成»request_tune()kernel_config()iter_kernels()built_entries(interface)held_delegates()对外接口__init__(*, target=None)forward(...)kernel_for(interface, call)key_for(interface, call)delegate_for(stage, identity, given=None, /, **params)roof_key()对内接口__init_subclass__() —— 类定义时运行，安装生成的方法_run_call(inputs, body)_bind_target()_reset_binding()_get_or_build_entry(interface, call)_check_signature(tensors)_infer_output_shapes(...) «生成»内部实现RMSNormFwdOpkernel_types : 3 个实现interfaces : rms_normforward(x, weight=None)MoEExpertMLPFwdOpdelegate_types = {gate_up, down}forward(expert_input, ...)_CompileBoundaryplan : _SignaturePlanbranches : effect 分支 → custom opcall(op, inputs, writes, execution)build_call_boundary(cls) : 子类的 _call_boundary()_SignaturePlansigconstructionschecksshapeseffectsroofscheck_construction(op) : indicescheck(op, tensors) : SignatureCalloutput_shapes(op, shapes, dtypes)effect_branch(op, present)eval_roofline(call)«dataclass, frozen»SignatureCallindicesbranch_keydevicetensorstrafficwrittenhas_outmetadatastages : stage → SignatureCall*metadata_values(name)Kerneldevicessupported_archsgeneralpreferred_overrefusal(call) [classmethod]entry_for(call) [classmethod]forward(...)KernelInterfacerequest : CallSpecforward(...)RMSNormFwdInterfacerequest = LayerNormCallforward(x, weight)entry运行时对象：缓存未命中时，由实现类 entry_for 返回的构造函数构造一个 kernel 或多个 kernel 的组合本例即一个 RMSNormOnChipKernel 实例«Op 子类»MoEGroupedGemmFwdOpforward(...)RMSNormOnChipKernelsupported_archs = [90]preferred_over = {rms_norm_streaming}refusal(call) [classmethod]20..1last_call0..1实现表[key]0..* +Op 类结构Opkernel_types : key → 实现类interfaces : interface → 接口类delegate_types : stage → Op 类injected_parameterscompile_op_names_signature : _SignaturePlan（同类实例共享） 类属性（每个 op 类一份）  实例状态 manifest 参数（构造期张量除外）target_installed_kernel_types : key → 实现类_preferred_over_keys_by_interface_construction_indices_instance_key构造_target_builder_serving_target_entries_by_call : (interface, call) → entry_built_entries : interface → entries_target_kernels : (设备, 各输入 dtype、shape) → kernel绑定构造期张量_tune_requested_delegates : stage → identity → Op_delegate_stages_effect_branches_last_call : SignatureCall_tune_warned运行 方法 __call__(*args) —— 参数即 forward 的签名serving_target «property»last_call «property»eval_roofline() «生成»request_tune()kernel_config()iter_kernels()_entries_built(interface)_held_delegates()对外接口__init__(*, target=None)forward(...)kernel_for(interface, call)key_for(interface, call)delegate_for(stage, identity, given=None, /, **params)roof_key()对内接口__init_subclass__() —— 类定义时运行，安装生成的方法_run_call(inputs, body)_bind_target()_reset_binding()_get_or_build_entry(interface, call)_check_signature(tensors)_infer_output_shapes(...) «生成»内部实现RMSNormFwdOpkernel_types : 3 个实现interfaces : rms_normforward(x, weight=None)MoEExpertMLPFwdOpdelegate_types = {gate_up, down}forward(expert_input, ...)_CompileBoundaryplan : _SignaturePlanbranches : effect 分支 → custom opcall(op, inputs, writes, execution)build_call_boundary(cls) : 子类的 _call_boundary()_SignaturePlansigconstructionschecksshapeseffectsroofscheck_construction(op) : indicescheck(op, tensors) : SignatureCalloutput_shapes(op, shapes, dtypes)effect_branch(op, present)eval_roofline(call)«dataclass, frozen»SignatureCallindicesbranch_keydevicetensorstrafficwrittenhas_outmetadatastages : stage → SignatureCall*metadata_values(name)Kerneldevicessupported_archsgeneralpreferred_overrefusal(call) [classmethod]entry_for(call) [classmethod]forward(...)KernelInterfacerequest : CallSpecforward(...)RMSNormFwdInterfacerequest = LayerNormCallforward(x, weight)entry运行时对象：缓存未命中时，由实现类 entry_for 返回的构造函数构造一个 kernel 或多个 kernel 的组合本例即一个 RMSNormOnChipKernel 实例«Op 子类»MoEGroupedGEMMFwdOpforward(...)RMSNormOnChipKernelsupported_archs = [90]preferred_over = {rms_norm_streaming}refusal(call) [classmethod]20..1last_call0..1实现表[key]0..* diff --git a/docs/user-guide/op/index.md b/docs/user-guide/op/index.md index 6323b61b..53db5cc4 100644 --- a/docs/user-guide/op/index.md +++ b/docs/user-guide/op/index.md @@ -178,7 +178,7 @@ The base class's members fall into three layers by who uses them: | 3 | `serving_target` | The target this instance selected | [Construction and calls § Targets and the implementation table](lifecycle.md#target) | | 4 | `last_call`, `eval_roofline()` | The last completed call, and its `(flops, bytes)` | [Call records and tuning § last_call](records.md#last-call) | | 5 | `request_tune()`, `kernel_config()` | Tuning, and reading the configuration in use | [Call records and tuning § Tuning](records.md#tune) | -| 6 | `iter_kernels()`, `built_entries(interface)`, `held_delegates()` | Enumerate the built kernels, entries and sub-ops | [Call records and tuning § Enumeration](records.md#enumerate) | +| 6 | `iter_kernels()`, `_entries_built(interface)`, `_held_delegates()` | Enumerate the built kernels, entries and sub-ops | [Call records and tuning § Enumeration](records.md#enumerate) | **Table 5** Internal interface: declared or called by op subclasses diff --git a/docs/user-guide/op/index.zh.md b/docs/user-guide/op/index.zh.md index 0206f962..6dbddbf1 100644 --- a/docs/user-guide/op/index.zh.md +++ b/docs/user-guide/op/index.zh.md @@ -139,7 +139,7 @@ op 实例构造一次，调用多次： | 3 | `serving_target` | 这个实例选定的 target | [构造与调用 § target 与实现表](lifecycle.md#target) | | 4 | `last_call`、`eval_roofline()` | 最近一次完成的调用，及其 `(flops, bytes)` | [调用记录与调优 § last_call](records.md#last-call) | | 5 | `request_tune()`、`kernel_config()` | 调优，及读取所用配置 | [调用记录与调优 § 调优](records.md#tune) | -| 6 | `iter_kernels()`、`built_entries(interface)`、`held_delegates()` | 枚举已构造的 kernel、entry 与子 op | [调用记录与调优 § 枚举](records.md#enumerate) | +| 6 | `iter_kernels()`、`_entries_built(interface)`、`_held_delegates()` | 枚举已构造的 kernel、entry 与子 op | [调用记录与调优 § 枚举](records.md#enumerate) | **表 5** 对内接口：op 子类声明或调用 diff --git a/docs/user-guide/op/lifecycle.md b/docs/user-guide/op/lifecycle.md index 952d4a69..07056c12 100644 --- a/docs/user-guide/op/lifecycle.md +++ b/docs/user-guide/op/lifecycle.md @@ -159,16 +159,9 @@ The signature check and dtypes: ## 4. Targets and the implementation table {#target} -A backend joins TileOPs in one of two ways; how to use them is in +A backend takes over a whole op through a target; how to use it is in [How a backend joins TileOPs](../dispatch/backends.md#choose). This section describes how the -base class handles them. - -**Table 4** The two ways a backend joins - -| No. | Way | What it takes over | When the base class handles it | -| --- | --- | --- | --- | -| 1 | `target` | the whole op | selected in step 4 of a call, and unchanged after that | -| 2 | `register_kernel_type` | adds one implementation to a kernel interface, which takes part in selection by its own declared applicability and precedence | merged into the implementation table by `Op.__init__` at construction | +base class handles it: the target is selected in step 4 of a call, and unchanged after that. ### Selecting the target {#select-target} @@ -262,7 +255,7 @@ This section distinguishes three words: One kernel type builds several entries, one per identity. Taking `RMSNormOnChipKernel` as the example: -**Table 5** Kernel type and entry +**Table 4** Kernel type and entry | No. | | Kernel type | Entry | | --- | --- | --- | --- | @@ -273,7 +266,6 @@ example: | 5 | Kept in | `kernel_types`, `_installed_kernel_types` | `_built_entries` | | 6 | Used for | the base class selects this call's implementation from it | the subclass calls it: `entry(x, weight)` | -`register_kernel_type` adds a kernel type to the candidates and builds no entry; `register_kernel_builder` registers a builder that builds the target's kernel directly, without selecting among kernel types; see [§ 4](#target). @@ -337,7 +329,7 @@ If the target was selected in this call's step 4: After that, the binding group is back to having no target selected; the records of completed calls and the held sub-ops are kept. The next call selects the target again. -**Table 6** Instance state after a failure +**Table 5** Instance state after a failure | No. | Case | Target | Built kernels | `last_call` | | --- | --- | --- | --- | --- | diff --git a/docs/user-guide/op/lifecycle.zh.md b/docs/user-guide/op/lifecycle.zh.md index c5094732..2740d383 100644 --- a/docs/user-guide/op/lifecycle.zh.md +++ b/docs/user-guide/op/lifecycle.zh.md @@ -112,14 +112,7 @@ ## 4. target 与实现表 {#target} -backend 接入 TileOPs 有两种方式，接入方法见 [backend 如何接入 TileOPs](../dispatch/backends.md#choose)。本节说明基类如何处理它们。 - -**表 4** backend 接入的两种方式 - -| No. | 方式 | 接管的范围 | 基类何时处理 | -| --- | --- | --- | --- | -| 1 | `target` | 整个 op | 调用的第 4 步选定，之后不再改变 | -| 2 | `register_kernel_type` | 为一个 kernel 接口新增一个实现，按它自己声明的适用范围与优先关系参与选择 | 构造时由 `Op.__init__` 并入实现表 | +backend 通过 target 接管整个 op，接入方法见 [backend 如何接入 TileOPs](../dispatch/backends.md#choose)。本节说明基类如何处理它：target 在调用的第 4 步选定，之后不再改变。 ### 选择 target {#select-target} @@ -192,7 +185,7 @@ target 被调用前，生成的检查已经运行，基类保证以下三点： 一个 kernel type 按 identity 构造出多个 entry。以 `RMSNormOnChipKernel` 为例： -**表 5** kernel type 与 entry +**表 4** kernel type 与 entry | No. | | kernel type | entry | | --- | --- | --- | --- | @@ -203,7 +196,7 @@ target 被调用前，生成的检查已经运行，基类保证以下三点： | 5 | 保存在 | `kernel_types`、`_installed_kernel_types` | `_built_entries` | | 6 | 用途 | 基类据此选择这次调用的实现 | 子类调用它：`entry(x, weight)` | -`register_kernel_type` 向候选中加入一个 kernel type，不构造任何 entry；`register_kernel_builder` 登记的 builder 直接构造 target 的 kernel，不经过 kernel type 的选择，见[第 4 节](#target)。 +`register_kernel_builder` 登记的 builder 直接构造 target 的 kernel，不经过 kernel type 的选择，见[第 4 节](#target)。 子类通过 `self.kernel_for(interface, call)` 取得 kernel： @@ -249,7 +242,7 @@ target 被调用前，生成的检查已经运行，基类保证以下三点： 处理之后，绑定组回到尚未选定 target 的状态；已完成调用的记录和持有的子 op 保留。下一次调用重新选择 target。 -**表 6** 失败后的实例状态 +**表 5** 失败后的实例状态 | No. | 情况 | target | 已构造的 kernel | `last_call` | | --- | --- | --- | --- | --- | diff --git a/docs/user-guide/op/records.md b/docs/user-guide/op/records.md index 67a49f2b..f895eff3 100644 --- a/docs/user-guide/op/records.md +++ b/docs/user-guide/op/records.md @@ -69,16 +69,16 @@ A subclass may also override two methods: | No. | Method | Returns | | --- | --- | --- | -| 1 | `iter_kernels()` | • Every `Kernel` in each interface's entries and in the held sub-ops, each once
• Kernels a target built are not walked; `built_entries(interface)` shows them | -| 2 | `built_entries(interface)` | • In-tree path: the entries built for this kernel interface, told apart by `(implementation class, identity)`
• When a target serves the op: every target kernel of the op, told apart by device and each input's dtype and shape; `interface` does not filter | -| 3 | `held_delegates()` | The sub-ops the composite op holds, in stage order | +| 1 | `iter_kernels()` | • Every `Kernel` in each interface's entries and in the held sub-ops, each once
• Kernels a target built are not walked; `_entries_built(interface)` shows them | +| 2 | `_entries_built(interface)` | • In-tree path: the entries built for this kernel interface, told apart by `(implementation class, identity)`
• When a target serves the op: every target kernel of the op, told apart by device and each input's dtype and shape; `interface` does not filter | +| 3 | `_held_delegates()` | The sub-ops the composite op holds, in stage order | What the three methods are for: -- `built_entries()` is for tests and benchmark reports only. +- `_entries_built()` is for tests and benchmark reports only. - `request_tune()` and `kernel_config()` reach the in-tree kernels through `iter_kernels()`, and the kernels a target built by device. -- `held_delegates()` also serves the recursive undo of sub-ops after a failure. +- `_held_delegates()` also serves the recursive undo of sub-ops after a failure. The call path does not get kernels through these three methods: @@ -111,7 +111,7 @@ reflection: kernel that does not define it makes the base class warn, once per instance. The rules are in [Adding a hardware backend § build_kernel](../../backends.md#build-kernel). - **Subclass overrides:** the above is the base class's behavior, and a subclass may override - it. `GemmW4A16FwdOp` does not support generic tuning: it overrides `request_tune()` to warn + it. `GEMMW4A16FwdOp` does not support generic tuning: it overrides `request_tune()` to warn and use its own calibrated selection. - **`kernel_config()`:** returns the op's own configuration; without one, the first non-empty `config` among the kernels `iter_kernels()` yields, then among the kernels a target built. diff --git a/docs/user-guide/op/records.zh.md b/docs/user-guide/op/records.zh.md index 580ea404..d21a5bbe 100644 --- a/docs/user-guide/op/records.zh.md +++ b/docs/user-guide/op/records.zh.md @@ -59,15 +59,15 @@ | No. | 方法 | 返回 | | --- | --- | --- | -| 1 | `iter_kernels()` | • 各接口的 entry 和持有的子 op 中的每个 `Kernel`，每个只返回一次
• 不遍历 target 构造的 kernel，它们由 `built_entries(interface)` 查看 | -| 2 | `built_entries(interface)` | • in-tree 路径：返回这个 kernel 接口已构造的 entry，按 `(实现类, identity)` 区分
• target 服务时：返回这个 op 的全部 target kernel，按设备与每个输入的 dtype、shape 区分，`interface` 不参与筛选 | -| 3 | `held_delegates()` | 组合 op 持有的子 op，按 stage 顺序 | +| 1 | `iter_kernels()` | • 各接口的 entry 和持有的子 op 中的每个 `Kernel`，每个只返回一次
• 不遍历 target 构造的 kernel，它们由 `_entries_built(interface)` 查看 | +| 2 | `_entries_built(interface)` | • in-tree 路径：返回这个 kernel 接口已构造的 entry，按 `(实现类, identity)` 区分
• target 服务时：返回这个 op 的全部 target kernel，按设备与每个输入的 dtype、shape 区分，`interface` 不参与筛选 | +| 3 | `_held_delegates()` | 组合 op 持有的子 op，按 stage 顺序 | 三个方法的用途： -- `built_entries()` 只用于测试与 benchmark 报告。 +- `_entries_built()` 只用于测试与 benchmark 报告。 - `request_tune()` 与 `kernel_config()` 经 `iter_kernels()` 取得 in-tree kernel，另按设备取得 target 构造的 kernel。 -- `held_delegates()` 还供失败时递归撤销子 op 使用。 +- `_held_delegates()` 还供失败时递归撤销子 op 使用。 调用路径不用这三个方法取得 kernel： @@ -87,5 +87,5 @@ - 对已构造的每个 kernel 调用 `Kernel.request_tune()`：程序已构建的 kernel 立即调优，尚未构建的在构建程序的那次启动时调优；没有 `autotune_configs` 的 kernel 不调优。 - 之后构造的 entry，基类构造后立即对其中每个 `Kernel` 调用 `request_tune()`。 - **target：** target 构造的 kernel 可选定义 `request_tune()`。`request_tune()` 对已构造的每个 target kernel，在它所属的每个设备上各调用一次，调用时该设备是当前设备；之后构造的 target kernel 在首次调用之前调用。kernel 没有定义它时发出警告，每个实例只警告一次。约定见[接入新的硬件 backend § build_kernel](../../backends.md#build-kernel)。 -- **子类覆盖：** 以上是基类的行为，子类可以覆盖。`GemmW4A16FwdOp` 不支持通用调优：覆盖 `request_tune()`，发出警告，改用自己校准过的选择。 +- **子类覆盖：** 以上是基类的行为，子类可以覆盖。`GEMMW4A16FwdOp` 不支持通用调优：覆盖 `request_tune()`，发出警告，改用自己校准过的选择。 - **`kernel_config()`：** 返回 op 自己的配置；没有时，先在 `iter_kernels()` 枚举的 kernel 中、再在 target 构造的 kernel 中，返回第一个非空的 `config`。 diff --git a/docs/user-guide/op/structure.md b/docs/user-guide/op/structure.md index 4991b86c..43ad0779 100644 --- a/docs/user-guide/op/structure.md +++ b/docs/user-guide/op/structure.md @@ -156,7 +156,7 @@ Notes on Table 2: | 1 | `_target_builder` | • The selected target's builder
• `None` means the in-tree kernels | class-attribute default `_UNRESOLVED`, meaning not selected yet | | 2 | `_serving_target` | The selected target's name, which the `serving_target` property reads | class-attribute default `None` | | 3 | `_entries_by_call` | The cache from `(interface, call)` to entry; a hit is one lookup | an empty table `Op.__init__` builds when it installs the implementation table | -| 4 | `_built_entries` | The entries under each interface, keyed by `(implementation class, identity)`; `built_entries` reads it | an empty table `Op.__init__` builds when it installs the implementation table | +| 4 | `_built_entries` | The entries under each interface, keyed by `(implementation class, identity)`; `_entries_built` reads it | an empty table `Op.__init__` builds when it installs the implementation table | | 5 | `_target_kernels` | • The kernels a target built, keyed by device and each input's dtype and shape
• An optional input that is not passed holds its own slot | an empty table `Op.__init__` builds | The base class's methods read and write these fields without asking whether they exist, diff --git a/docs/user-guide/op/structure.zh.md b/docs/user-guide/op/structure.zh.md index 9a8ef3aa..b087e462 100644 --- a/docs/user-guide/op/structure.zh.md +++ b/docs/user-guide/op/structure.zh.md @@ -95,7 +95,7 @@ | 1 | `_target_builder` | • 选定 target 的 builder
• `None` 表示 in-tree kernel | 类属性默认值 `_UNRESOLVED`，表示尚未选定 | | 2 | `_serving_target` | 选定的 target 名，`serving_target` 属性读取它 | 类属性默认值 `None` | | 3 | `_entries_by_call` | `(interface, call)` 到 entry 的缓存，命中只需一次查表 | `Op.__init__` 安装实现表时建立的空表 | -| 4 | `_built_entries` | 每个接口下按 `(实现类, identity)` 保存的 entry，`built_entries` 读取它 | `Op.__init__` 安装实现表时建立的空表 | +| 4 | `_built_entries` | 每个接口下按 `(实现类, identity)` 保存的 entry，`_entries_built` 读取它 | `Op.__init__` 安装实现表时建立的空表 | | 5 | `_target_kernels` | • target 构造的 kernel，按设备与每个输入的 dtype、shape 保存
• 未传入的可选输入单独占位 | `Op.__init__` 建立的空表 | 基类的方法读写这些字段时，不判断字段是否存在，因为构造之后它们都存在： diff --git a/figures/user-guide/dispatch/extension.en.puml b/figures/user-guide/dispatch/extension.en.puml index 6a2198f1..572b3092 100644 --- a/figures/user-guide/dispatch/extension.en.puml +++ b/figures/user-guide/dispatch/extension.en.puml @@ -35,8 +35,6 @@ package "The op instance's in-tree implementation" { rectangle "The Op class's forward" as body <> package "Implementations of one kernel interface\n(installed at construction)" as snap { rectangle "Classes in kernel_types\ninheriting the interface" as intree <> - rectangle "register_kernel_type\nadds a key\n(every instance built afterwards)" as added <> - intree -[hidden]down- added } } diff --git a/figures/user-guide/dispatch/extension.zh.puml b/figures/user-guide/dispatch/extension.zh.puml index fba7f8c3..622e9942 100644 --- a/figures/user-guide/dispatch/extension.zh.puml +++ b/figures/user-guide/dispatch/extension.zh.puml @@ -35,8 +35,6 @@ package "op 实例的 in-tree 实现" { rectangle "Op 类的 forward" as body <> package "一个 kernel 接口的实现（构造时安装）" as snap { rectangle "kernel_types 中\n继承该接口的类" as intree <> - rectangle "register_kernel_type\n新增的 key\n（此后构造的所有实例）" as added <> - intree -[hidden]down- added } } diff --git a/figures/user-guide/op/structure.en.puml b/figures/user-guide/op/structure.en.puml index 1210e5f4..22f57d0c 100644 --- a/figures/user-guide/op/structure.en.puml +++ b/figures/user-guide/op/structure.en.puml @@ -65,8 +65,8 @@ abstract class Op #E9EAF0 { request_tune() kernel_config() iter_kernels() - built_entries(interface) - held_delegates() + _entries_built(interface) + _held_delegates() .. internal interface .. __init__(*, target=None) {abstract} forward(...) @@ -164,7 +164,7 @@ object "entry" as Entry #EEF2E6 { here, one RMSNormOnChipKernel instance } -class MoEGroupedGemmFwdOp <> #E9EAF0 { +class MoEGroupedGEMMFwdOp <> #E9EAF0 { forward(...) } @@ -178,7 +178,7 @@ class RMSNormOnChipKernel #EEF2E6 { ' subclasses on top RMSNormFwdOp -down-|> Op MoEExpertMLPFwdOp -down-|> Op -MoEExpertMLPFwdOp o-right-> "2" MoEGroupedGemmFwdOp +MoEExpertMLPFwdOp o-right-> "2" MoEGroupedGEMMFwdOp ' code generation below Op .left.> _CompileBoundary diff --git a/figures/user-guide/op/structure.zh.puml b/figures/user-guide/op/structure.zh.puml index feb959bd..3c5ebc75 100644 --- a/figures/user-guide/op/structure.zh.puml +++ b/figures/user-guide/op/structure.zh.puml @@ -65,8 +65,8 @@ abstract class Op #E9EAF0 { request_tune() kernel_config() iter_kernels() - built_entries(interface) - held_delegates() + _entries_built(interface) + _held_delegates() .. 对内接口 .. __init__(*, target=None) {abstract} forward(...) @@ -164,7 +164,7 @@ object "entry" as Entry #EEF2E6 { 本例即一个 RMSNormOnChipKernel 实例 } -class MoEGroupedGemmFwdOp <> #E9EAF0 { +class MoEGroupedGEMMFwdOp <> #E9EAF0 { forward(...) } @@ -178,7 +178,7 @@ class RMSNormOnChipKernel #EEF2E6 { ' 子类在上 RMSNormFwdOp -down-|> Op MoEExpertMLPFwdOp -down-|> Op -MoEExpertMLPFwdOp o-right-> "2" MoEGroupedGemmFwdOp +MoEExpertMLPFwdOp o-right-> "2" MoEGroupedGEMMFwdOp ' 代码生成在下 Op .left.> _CompileBoundary diff --git a/hooks.py b/hooks.py index 92ce0f25..6169eeb1 100644 --- a/hooks.py +++ b/hooks.py @@ -1,9 +1,10 @@ """mkdocs hooks: mirrored-content path rewrites, the Benchmarks nav, and the untranslated-page notice. -* Mirrored design docs carry repo-relative paths (`../../tileops/...`) that do - not resolve here; `on_page_markdown` rewrites them after include-markdown has - pulled the content in. +* Mirrored design docs carry repo-relative paths (`../../tileops/...`) and + links to other upstream docs (`../code-style.md`) that do not resolve here; + `on_page_markdown` rewrites them after include-markdown has pulled the content + in. * The Benchmarks pages are generated at deploy time and cannot be listed in `mkdocs.yml`; `on_config` expands that one nav entry to whatever the renderer produced. @@ -32,6 +33,7 @@ UPSTREAM_BLOB = "https://github.com/tile-ai/TileOPs/blob/main" _DESIGN_REPO_PATH = re.compile(r"\.\./\.\./([\w./-]+)") +_DESIGN_DOCS_PATH = re.compile(r"\]\(\.\./([\w./-]+)") # Keyed by the locale being built, not by the locale of the content shown. # A plain block, not an admonition: admonitions are styled as technical asides @@ -65,6 +67,8 @@ def on_page_markdown(markdown, page, config, files): # Upstream design docs link to source files via ../../. # Redirect those to GitHub so they resolve from the published site. markdown = _DESIGN_REPO_PATH.sub(rf"{UPSTREAM_BLOB}/\1", markdown) + # And to upstream docs outside docs/design via ../. + markdown = _DESIGN_DOCS_PATH.sub(rf"]({UPSTREAM_BLOB}/docs/\1", markdown) notice = _fallback_notice(page) if notice: @@ -245,6 +249,7 @@ def on_page_content(html, page, config, files): # redirect is written into each locale's build. _REDIRECTS = { "api/topk/": "../sampling/#top-k-selection", + "design/architecture/": "../", } _REDIRECT_PAGE = ( diff --git a/mkdocs.yml b/mkdocs.yml index 21f0c0b9..45242744 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -106,7 +106,6 @@ nav: # an op against it, then how an op is accepted. - Design: - design/index.md - - Architecture: design/architecture.md - Layer Boundaries: design/layer-boundaries.md - Manifest: design/manifest.md - Op Interfaces: design/ops-design.md @@ -204,7 +203,6 @@ plugins: Writing Benchmarks: 编写 benchmark Using torch.compile: 接入 torch.compile Adding a Backend: 接入新硬件 backend - Architecture: 架构 Manifest: manifest 规范 Op Interfaces: op 接口 Op Interface Reference: op 接口参考 diff --git a/scripts/check_api_pages.py b/scripts/check_api_pages.py index a6f88ef9..921347b3 100644 --- a/scripts/check_api_pages.py +++ b/scripts/check_api_pages.py @@ -10,7 +10,7 @@ Families come from `_FAMILIES` in `tileops/__init__.py`, their ops from the `__all__` of `tileops/.py`, read with `ast`: importing a family pulls in torch, which the docs environment does not install. A path of more than two -segments, such as `tileops.trace.api._Trace`, belongs to no family and is left +segments, such as `tileops.trace._api._Trace`, belongs to no family and is left alone. Usage: diff --git a/scripts/gen_bench_pages.py b/scripts/gen_bench_pages.py index b7f0845e..a7c9713b 100644 --- a/scripts/gen_bench_pages.py +++ b/scripts/gen_bench_pages.py @@ -218,7 +218,7 @@ def family_of(op: str, op_module: str | None, return fam mod = (op_module or "").lower() parts = mod.split(".") - if len(parts) >= 4 and parts[0] == "tileops" and parts[1] == "ops": + if len(parts) >= 4 and parts[0] == "tileops" and parts[1] in ("ops", "_ops"): if parts[2] in _MODULE_FAMILY: return _MODULE_FAMILY[parts[2]] hay = f"{mod} {op.lower()}" @@ -364,7 +364,7 @@ def load_sol_engine(gpu: str, tileops: str = TILEOPS): os.path.join(tileops, "scripts", "nightly_report.py")) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) - from tileops.perf.profile import find_profile + from tileops._perf import find_profile profile = find_profile(gpu) except Exception as exc: # noqa: BLE001 — degrade to an empty column print(f"warning: SOL column disabled ({exc})", file=sys.stderr) diff --git a/scripts/workload_shape.py b/scripts/workload_shape.py index bde76517..4076f7c6 100644 --- a/scripts/workload_shape.py +++ b/scripts/workload_shape.py @@ -393,9 +393,8 @@ def _load(self): src = os.path.join(self.tileops, "src") if src not in sys.path: sys.path.insert(0, src) - from tileops.manifest import workload - from tileops.manifest.plan import entry_plan - self._tools = (entry_plan, workload) + from tileops.manifest import entry_plan, instantiate + self._tools = (entry_plan, instantiate) except Exception as exc: # noqa: BLE001 — degrade to ids alone print(f"warning: parametric workloads keep their ids alone " f"(tileops.manifest not importable: {exc})", file=sys.stderr) @@ -406,7 +405,7 @@ def describe(self, op: str, entry: dict, config: str) -> Spec | None: tools = self._load() if not tools: return None - entry_plan, workload = tools + entry_plan, instantiate = tools try: if op not in self._plans: self._plans[op] = entry_plan(op, entry, self.adts, resolve=False) @@ -419,11 +418,11 @@ def describe(self, op: str, entry: dict, config: str) -> Spec | None: continue for case in row.get("dtype_cases") or [{}]: try: - call = workload.instantiate(plan, row, case) + call = instantiate(plan, row, case) except Exception: # noqa: BLE001 — a row this checkout rejects continue if call.case_id == config: - return _parametric_spec(plan, row, case, call, workload) + return _parametric_spec(plan, row, case, call) return None @@ -438,13 +437,15 @@ def _fmt_param(value) -> str: return _fmt_value(value) -def _symbolic_branch(plan, row, call, workload): +def _symbolic_branch(plan, row, call): """Each present input's shape in the manifest's symbols, type families expanded. None when the branch cannot be rebuilt, so the row prints concrete shapes. """ sig = plan.sig try: + # `tileops.manifest` exports no way to pick a row's branch. + from tileops.manifest import _workload as workload scope = {p: workload._convert(sig, p, row[p] if p in row else sig.params[p].get("default")) for p in sig.params} @@ -463,7 +464,7 @@ def _symbolic_branch(plan, row, call, workload): return out -def _parametric_spec(plan, row, case, call, workload) -> Spec: +def _parametric_spec(plan, row, case, call) -> Spec: sig = plan.sig present = [(n, call.specs[n]) for n in sig.inputs if call.specs.get(n) is not None] # The row's dtype is the first one its case id names, so a tensor in @@ -479,7 +480,7 @@ def _parametric_spec(plan, row, case, call, workload) -> Spec: tensors = _group_tensors((n, fmt_shape(d), dt) for n, d, dt in shapes) symbolic, bindings = None, OrderedDict() - sym = _symbolic_branch(plan, row, call, workload) + sym = _symbolic_branch(plan, row, call) if sym is not None and len(sym) == len(present): for text in sym.values(): for node in ast.walk(ast.parse(text, mode="eval")): diff --git a/tests/test_renderer.py b/tests/test_renderer.py index e64e5bef..051ef9b0 100644 --- a/tests/test_renderer.py +++ b/tests/test_renderer.py @@ -151,6 +151,7 @@ def test_the_package_decides_the_family_not_a_word_in_the_name(): assert g.family_of("DeltaDecodeFwdOp", "tileops.ops.linear_attention.delta") == "linear_attention" assert g.family_of("GemmOp", "tileops.ops.gemm.gemm") == "linear_algebra" + assert g.family_of("GEMMFwdOp", "tileops._ops.gemm.gemm") == "linear_algebra" # An op defined in a module rather than a package still falls through. assert g.family_of("RmsNormFwdOp", None) == "normalization"