跳到主要内容
TileLang 系统教程

tilelang-tutorial.huandu.me

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

/code/28-operator-design-methodology/

06-case-norm-rope-quant/README.md2821 字节
# 示例 06:完整案例——QK-Norm + RoPE + fp8 量化

## 目的

这是第 28 章第 6 节案例的可复算版本:同一个算子,三种设计,
`预算表 → 第一版 → 性能诊断 → 第二版`。

被建模的算子(推理侧 attention 前处理):

```
q,k = proj(x)                         # bf16, [tokens, heads, head_dim]
q,k = rmsnorm(q,k) per (token, head)  # 沿 head_dim
q,k = rope(q,k)                       # 每个 token 所有 head 共用同一张表
q8,k8, scale = fp8_quant(q,k)         # 每行一个 fp32 scale
```

三种设计:

| 设计 | 一个 block 处理 | 特点 |
| --- | --- | --- |
| A | 一行(一个 (token, head) 的 128 个元素) | 最直觉的写法:一行一 block |
| B | 连续的 8 行(对齐到 head 组) | 把查表摊到 8 行上 |
| C | 一整个 token(全部 q 和 k 的 40 行) | 把查表摊到整个 token |

模型只数两件事:**payload**(输入、输出、scale)和 **lookup**(norm 权重、RoPE
cos/sin 表)。lookup 的关键不是它有多大,而是**它被重复读了多少次**。

## 运行前提

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

## 运行命令

```bash
python3 case_norm_rope_quant.py
python3 case_norm_rope_quant.py --rows-per-block 8
python3 case_norm_rope_quant.py --tokens 8192 --bandwidth-gbs 3350
python3 case_norm_rope_quant.py --q-heads 64 --kv-heads 8   # 别的模型配置
```

## 期望输出

正确性判据:退出码为 0。默认参数下的实际输出(节选):

```text
part 0 -- the operator ledger
flops                   : 210.2 MFLOP
ideal bytes             : 67.76 MB   (payload 63.57 MB + lookup 4.19 MB)
arithmetic intensity    : 3.102 FLOP/Byte
memory roof at 4502.4 GB/s : 15.051 us

part 1 -- budget table
design                     rows/blk    grid  waves  on-chip KB  payload MB  lookup MB  total MB  vs ideal
A: one row per block              1  163840    156         1.9       63.57     251.66    315.23     4.65x
B: 8 rows per block               8   20480     20         8.0       63.57      31.46     95.03     1.40x
C: one whole token per block     40    4096      4        36.2       63.57       6.29     69.86     1.03x

part 2 -- diagnosis
A: one row per block
    total  315.23 MB = payload  63.57 MB + lookup  251.66 MB   (lookup is  79.8% of all traffic)
    over the ideal by  247.46 MB  -> the achievable time floor rises from 15.05 us to 70.01 us
B: 8 rows per block   -> 95.03 MB, floor 21.11 us
C: one whole token    -> 69.86 MB, floor 15.52 us
```

这组数字就是案例的骨架:**第一版(A)比理想多搬 4.65 倍的字节,其中 79.8% 是查表;
把粒度和 head 组对齐之后(B/C),剩下的差距只有 3%。**

## 性能口径

**不测量任何性能。** 只有 `bytes / bandwidth` 一类的下界。案例里凡是要说
「实测多少」的地方,正文都写成「在有卡机器上应该这样验证」,不给数字。

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

**CPU 逻辑验证。** 在 macOS (Apple Silicon)、Python 3.11 上运行过默认、
`--rows-per-block`、`--tokens`、`--q-heads/--kv-heads` 多组参数;输出与本文件一致。
没有在 GPU/NPU 上运行过任何内容。

## 已知局限

* **lookup 的重复次数是设计意图,不是实测的 cache 行为。** 真实机器上 L2 完全可能
  把权重和表兜住,那时 A 设计的实际流量会接近 B/C。脚本把两种情况分开说:
  「模型上被重复读了几次」由设计决定,「实际少读了多少」由 cache 决定。
  `../02-roofline-misuse/` 的 `M6` 是同一件事的另一个侧面。
* 片上占用只算「一个 block 同时要活着的 buffer」,不含流水多版本。
* 没有建模 DMA 描述符开销:A 设计有 163840 个 block,每个 block 只搬 256 B 输入,
  这个开销在真实机器上只会让 A 更差,所以模型对 A 是**乐观**的。
* 波数用一个假设的 `blocks_per_SM`(默认 8)算,真实占用要用 `../04-tile-size-decision/`
  去查。

← 浏览本章全部代码