跳到主要内容
TileLang 系统教程

tilelang-tutorial.huandu.me

第 28 章 算子设计方法论与性能模型

/code/28-operator-design-methodology/

03-carver-cost-model/README.md3545 字节
# 示例 03:复算 carver/roller 的 tile cost model

## 目的

`_vendor/tilelang-v0.1.15/tilelang/carver/` 是 TileLang 自带的 tile 配置推荐器。
它的 cost model 决定了哪些 tile 会被推荐、哪些会被直接丢弃。这个脚本把那套算术
用标准库重写一遍,好在没有 GPU 的机器上读懂它:

| 量 | carver 里的位置 |
| --- | --- |
| `coalesced_factor` / `coalesced_tensor_shape` | `carver/roller/policy/common.py:32-52` |
| traffic | `carver/roller/policy/default.py:398-432` |
| shared memory(best-fit,32 B 对齐) | `carver/roller/policy/default.py:452-495` + `roller/bestfit.py:22-61` |
| `reg_usage = 2 * max(prod(tile) * bits / 32)`、容量上限 | `carver/roller/policy/default.py:556-574` |
| `block_per_SM`、`num_wave` | `carver/roller/policy/default.py:568-573` |
| 排序 `priority = (traffic + 1) * num_wave` | `carver/roller/policy/default.py:111-112` |
| 设备模型(`smem_cap` / `reg_cap` / `sm_partition` / `transaction_size`) | `carver/arch/cuda.py:153-167`、`arch/arch_base.py:6-27` |

`--self-test` 会把复算的 traffic 与教科书公式
`b·M·N·K·(1/block_N + 1/block_M) + b·M·N` 逐项对比。

## 运行前提

* Python 3.9+,只用标准库。
* **不需要 GPU / NPU / TileLang / PyTorch。**

注意:**不能**在无卡机器上 `import tilelang.carver` 来对照。`CUDA.__init__` 会调用
`tvm.runtime.cuda(0)` 并在没有设备时抛
`RuntimeError("Cannot find cuda device 0.")`(`carver/arch/cuda.py:144-146`),
`carver/arch/__init__.py:47-56` 的 `auto_infer_current_arch()` 在没有 CUDA 时会退到
CPU 模型。所以这份复算是**代码阅读的产物**,只保证算术与源码一致,不保证与运行时
完全一致。

## 运行命令

```bash
python3 carver_cost_model.py --self-test --topk 8
python3 carver_cost_model.py --m 8192 --n 8192 --k 8192 --topk 12
python3 carver_cost_model.py --smem-cap 101376 --cores 108 --topk 8   # 别的卡
```

## 期望输出

正确性判据:退出码为 0,且 `--self-test` 的 `rel err` 全为 `0.00e+00`。

```text
self-test: carver-style traffic vs the closed form b*M*N*K*(1/block_N + 1/block_M) + b*M*N
 block_M block_N block_K    model MB   closed MB   rel err
      64      64      64      4328.5      4328.5  0.00e+00
     128     128      32      2181.0      2181.0  0.00e+00
     128     256      64      1644.2      1644.2  0.00e+00

candidates visited : 729   valid : 154
rejected on regs     :  122
rejected on smem     :  453
model priority values in use : 48
ties at the top     : 8

  #  block_M block_N block_K  traffic MB  smem KB blk/SM waves        prio  binding
  1      128     256      16      1644.2     12.0      1     4    6.58e+09  regs
  ...
```

三个可以直接读出来的结论:

1. **traffic 对 `block_K` 不敏感**:`128×256×16` 与 `128×256×128` 的 traffic 完全
   相同,`prio` 也完全相同。模型的这一维没有分辨率。
2. **寄存器比 shared memory 先撞墙**:122 个候选被 `reg_usage > reg_cap` 丢掉,
   其中最大的一批是 `block_K` 很大而 `M×N` 不小的形状。
3. **模型会给出大量并列**:154 个可行候选只产出 48 个不同的 `prio` 值,前 8 名并列。
   并列就是「接下来该去测量」的信号(第 11 章 3.9 节的说法,这里给出了它的量级)。

## 性能口径

**不测量任何性能。** `prio` 是无量纲的排序键,不要当成时间。

## 作者实际做到的验证级别

**CPU 逻辑验证。** 在 macOS (Apple Silicon)、Python 3.11 上运行过默认、
`--self-test`、多组 `--m/--n/--k` 与多组 `--smem-cap/--cores`;输出与本文件一致。
没有在 GPU 上运行过 carver 本体。

## 已知局限

* **tile 传播是手写的。** carver 走的是 `PrimFuncNode.propagate`
  (`roller/node.py:332-399`)这一套通用的区域推导;本脚本对普通 GEMM 把它解析地
  写成「输出 tile 需要 A 的 `block_m` 行整条 K、B 的 `block_n` 行整条 K」。
  这一读法能精确复现教科书 traffic 公式;另一种读法(只取一个 `rstep` 切片)
  不能。要判定 carver 到底返回哪一个,需要在有卡的机器上实际跑一次——本机做不到,
  所以这里把它写成显式假设而不是结论。
* 候选枚举是一次性的因子笛卡尔积,不是 carver 的 `dfs_smem_tile` 最优先搜索
  (`roller/policy/default.py:96-134`,上限 2000 个已访问 tile)。在同样的候选集合
  上排序键一致,但访问顺序不同。
* `smem_cap` / `cores` 由命令行给出,不是从设备查询的。
* `reg_usage` 里的位宽取的是**累加器**的位宽(fp32 时 32)。carver 用的是
  `node.get_dtype()`,它在我们的复算里被取成累加 dtype;如果真实节点 dtype
  是输入 dtype(fp16 时 16),`reg_usage` 会小一倍,`block_per_SM` 的寄存器项
  会翻倍。要确认这一点同样需要在有卡的机器上跑一次 carver。
* 只建模普通 GEMM;split-K、grouped、epilogue 都不在模型里。

← 浏览本章全部代码